numpy：有条件的总和

问题内容：

我有以下numpy数组：

import numpy as np
arr = np.array([[1,2,3,4,2000],
                [5,6,7,8,2000],
                [9,0,1,2,2001],
                [3,4,5,6,2001],
                [7,8,9,0,2002],
                [1,2,3,4,2002],
                [5,6,7,8,2003],
                [9,0,1,2,2003]
              ])

我知道np.sum(arr, axis=0)可以提供以下结果：

array([   40,    28,    36,    34, 16012])

我想做的（ 没有for循环 ）是根据最后一列的值对各列求和，以便提供的结果是：

array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

我意识到这可能是一个无循环的尝试，但希望能做到最好……

如果必须使用for循环，那将如何工作？

我尝试了np.sum(arr[:, 4]==2000, axis=0)（2000用for循环中的变量替代），但是结果为 2

问题答案：

您可以使用一个巧妙的应用在纯numpy的做到这一点np.diff和np.add.reduceat。np.diff将为您提供最右边一列更改的索引：

d = np.diff(arr[:, -1])

np.where会将您的布尔索引d转换为np.add.reduceat期望的整数索引：

d = np.where(d)[0]

reduceat 还会期望看到零索引，并且所有内容都需要移动一个：

indices = np.r_[0, e + 1]

np.r_比起np.concatenate使用标量，使用此处要方便得多。然后，总和变为：

result = np.add.reduceat(arr, indices, axis=0)

当然，这可以组合成一种形式：

>>> result = np.add.reduceat(arr, np.r_[0, np.where(np.diff(arr[:, -1]))[0] + 1], axis=0)
>>> result
array([[   6,    8,   10,   12, 4000],
       [  12,    4,    6,    8, 4002],
       [   8,   10,   12,    4, 4004],
       [  14,    6,    8,   10, 4006]])

numpy：有条件的总和

微信关注