矩阵每行三个元素之和不带循环

Question

我想写一个函数SumThree，它接收一个大小为(n x (n+2))的矩阵作为numpy数组，并返回一个向量v，它的定义如下：v[i] = A[i,i] + A[i,i+1] + A[i,i+2]。你知道吗

限制：没有循环，您不能使用函数sum（np.sum）

我确实通过使用递归和更改函数的签名（函数应该只接收矩阵）来解决这个问题。但我在寻找一个更简单更有效的解决方案。你知道吗

像这样的事情（这就是我尝试的）：

def sumThree(array, curr=0, i=0, v=[]):
 if curr < array.size
   curr += array.shape[1]
 else
   return v
 count = array.size
 v[i] = array[curr,curr] + array[curr,curr+1] + array[curr,curr+2]
 sumThree(array, curr, i+1, v)

我还没试过，但我认为这应该管用：

def sumThree(a):
  return (a.diagonal() + a.diagonal(1) + a.diagonal(2))

Answer 1

您可以使用numpy.diagonal获得特定的对角线（具有所需的偏移量）

你能试试这个，看看对你有用吗？你知道吗

def sumThree(a):
    return a.diagonal(offset=0) + a.diagonal(offset=1) + a.diagonal(offset=2)

下面是一个完整的示例：

import numpy as np

def sumThree(a):
    return a.diagonal(offset=0) + a.diagonal(offset=1) + a.diagonal(offset=2)


my_arr = np.arange(48).reshape(6,8)
print(my_arr)
print("\n\n\n")
print(sumThree(my_arr))

它输出：

[[ 0  1  2  3  4  5  6  7]
 [ 8  9 10 11 12 13 14 15]
 [16 17 18 19 20 21 22 23]
 [24 25 26 27 28 29 30 31]
 [32 33 34 35 36 37 38 39]
 [40 41 42 43 44 45 46 47]]




[  3  30  57  84 111 138]