是否有函数计算返回一维数组结果的numpy ndarray元素的总和?

简单的例子:

a = array([[[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]],

       [[1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]]])

result = []

for i in a:
    result.append(i.sum())

result = [6, 3]

是否有一个numpy函数这样做更快?如果有帮助,则:a仅包含对角矩阵。

评论
  • 蔓延2j
    蔓延2j 回复

    A proper use of the axis parameter of np.sum() would do:

    import numpy as np
    
    
    a = np.array(
        [[[1, 0, 0],
          [0, 2, 0],
          [0, 0, 3]],
         [[1, 0, 0],
          [0, 1, 0],
          [0, 0, 1]]])
    
    
    np.sum(a, axis=(1, 2))
    # [6, 3]