在二维数组中沿对角线和反对角线线求和 - NumPy / Python

Question

我有一个如下的numpy数组：

np.array([[1,0,0],[1,0,0],[0,1,0]])
output:
array([[1, 0, 0],
       [1, 0, 0],
       [0, 1, 0]])

我希望将左对角线和右对角线的元素相加，生成一个新数组：

1）左对角线： 输出结果：

[1,1,0,1,0]

2) 右对角线：

输出结果：

[0,0,1,2,0]

有没有简便的方法？谢谢~

- james

Answer 1

方法1：使用掩码

下面是一个基于“掩码”的向量化方法 -

def left_diag_sum_masking(a):
    n = len(a)
    N = 2*n-1

    R = np.arange(N)
    r = np.arange(n)

    mask = (r[:,None] <= R) & (r[:,None]+n > R)

    b = np.zeros(mask.shape,dtype=a.dtype)
    b[mask] = a.ravel()

    return b.sum(0)

所以，`left_diag_sum_masking` 为我们提供了左对角线的总和。要获取右对角线的总和，只需沿列翻转、求和，然后再翻转回来即可。

因此，只需执行 -

right_diag_sum = left_diag_sum_masking(a[::-1])[::-1]

样例运行 -

In [220]: np.random.seed(0)

In [221]: a = np.random.randint(0,9,(4,4))

In [222]: a
Out[222]: 
array([[5, 0, 3, 3],
       [7, 3, 5, 2],
       [4, 7, 6, 8],
       [8, 1, 6, 7]])

In [223]: left_diag_sum_masking(a)
Out[223]: array([ 5,  7, 10, 23,  9, 14,  7])

In [224]: left_diag_sum_masking(a[::-1])[::-1] # right-diag sums
Out[224]: array([ 3,  5, 13, 21, 20,  5,  8])

方法 #2：使用零填充

def left_diag_sum_zerospad(a):
    n = len(a)
    N = 2*n-1
    p = np.zeros((n,n),dtype=a.dtype)
    ap = np.concatenate((a,p),axis=1)
    return ap.ravel()[:n*N].reshape(n,-1).sum(0)

获取右对角线之和-

right_diag_sum = left_diag_sum_zerospad(a[::-1])[::-1]

- Divakar

---

你的解决方案对于非方阵矩阵无效。请检查我的答案，其中包含适用于任意情况的代码。 - Andrew Sonin

Answer 2

另一种基于填充和滚动的方法：

def sum_shifted(arr, direction=1):
    n = arr.shape[0]
    temp = np.zeros((n, 2 * n - 1), dtype=arr.dtype)
    temp[:, slice(None, n) if direction == 1 else slice(-n, None)] = arr
    for i in range(n):
        temp[i, :] = np.roll(temp[i, :], direction * i)
    return np.sum(temp, 0)[::direction]

这为您提供了很多选择。就速度而言，@Divakar的方法似乎更有优势：

这些图表是使用此脚本生成的，使用以下函数作为测试函数：

def sum_shifted(arr, direction=1):
    n = arr.shape[0]
    temp = np.zeros((n, 2 * n - 1), dtype=arr.dtype)
    temp[:, slice(None, n) if direction == 1 else slice(-n, None)] = arr
    for i in range(n):
        temp[i, :] = np.roll(temp[i, :], direction * i)
    return np.sum(temp, 0)[::direction]


def sum_shifted_both(arr):
    return sum_shifted(arr, 1), sum_shifted(arr, -1)


def sum_adam(arr):
    return (
        np.array([np.sum(np.diag(np.fliplr(arr), d)) for d in range(len(arr) - 1, -len(arr), -1)]),
        np.array([np.sum(np.diag(arr, d)) for d in range(len(arr) - 1, -len(arr), -1)]))


def sum_divakar(a):
    n = len(a)
    N = 2*n-1

    R = np.arange(N)
    r = np.arange(n)

    mask = (r[:,None] <= R) & (r[:,None]+n > R)

    b_leftdiag = np.zeros(mask.shape,dtype=a.dtype)
    b_leftdiag[mask] = a.ravel()

    b_rightdiag = np.zeros(mask.shape,dtype=a.dtype)
    b_rightdiag[mask[:,::-1]] = a.ravel()

    return b_leftdiag.sum(0), b_rightdiag.sum(0)[::-1]


def sum_divakar2(a):
    def left_sum(a):
        n = len(a)
        N = 2*n-1
        p = np.zeros((n,n),dtype=a.dtype)
        ap = np.concatenate((a,p),axis=1)
        return ap.ravel()[:n*N].reshape(n,-1).sum(0)

    return left_sum(a), left_sum(a[::-1])[::-1]

并作为辅助函数：

def gen_input(n):
    return np.arange(n * n).reshape((n, n))


def equal_output(out_a, out_b):
    return all(
        np.all(a_arr == b_arr)
        for a_arr, b_arr in zip(out_a, out_b))


input_sizes=(5, 10, 50, 100, 500, 1000, 5000)
funcs = sum_shifted_both, sum_adam, sum_divakar, sum_divakar2


runtimes, input_sizes, labels, results = benchmark(
    funcs, gen_input=gen_input, equal_output=equal_output, input_sizes=input_sizes)

plot_benchmarks(runtimes, input_sizes, labels)

- norok2

Answer 3

你可以像这样使用.diag：

import numpy as np

arr = np.array([[1, 0, 0], [1, 0, 0], [0, 1, 0]])

left_diag = [np.sum(np.diag(np.fliplr(arr), d)) for d in range(len(arr) - 1, -len(arr), -1)]
right_diag = [np.sum(np.diag(arr, d)) for d in range(len(arr) - 1, -len(arr), -1)]

print("left:", left_diag)
print("right:", right_diag)

输出：

left: [1, 1, 0, 1, 0]
right: [0, 0, 1, 2, 0]

它返回给定偏移量的对角线上的所有元素。要按你提到的顺序获取所有偏移量，我们从+2到-2，然后对于每个对角线获取元素的总和。

要获取左对角线，我们首先使用.fliplr翻转 arr。

"Original Answer"的翻译是“最初的回答”。

- Adam.Er8

Answer 4

一个快速的解决方案，即使对于非方阵也有效：

from numpy.lib.stride_tricks import as_strided


def left_diag_sum(matrix):

    n_rows, n_cols = matrix.shape
    if n_rows > n_cols:
        matrix = matrix.T
        n_rows, n_cols = n_cols, n_rows
    diag_len = n_rows
    n_diags  = n_rows + n_cols - 1

    dtype   = matrix.dtype
    leaning = np.zeros((diag_len, n_diags), dtype)

    col_dim_stride = dtype.itemsize
    row_dim_stride = col_dim_stride * (n_diags + 1)

    strided      = as_strided(leaning, matrix.shape, (row_dim_stride, col_dim_stride))
    strided[...] = matrix

    return leaning.sum(0)


def right_diag_sum(matrix):

    row_reversed_matrix = matrix[::-1]
    return left_diag_sum(row_reversed_matrix)[::-1]

- Andrew Sonin

Answer 5

numpy是一个非常有用的函数，用于计算迹。以下是一个例子。

''' import numpy as np a=np.array([[1,0,0],[1,0,0],[0,1,0]]) b=np.trace(a) print(b) ''' 这里是文档链接

- Ajit