如何通过id高效地对2D NumPy数组进行求和和平均？

Question

我有一个二维数组a和一个一维数组b。 我想通过b中的每个id分组计算数组a中的行之和。 例如：

import numpy as np

a = np.array([[1,2,3],[2,3,4],[4,5,6]])
b = np.array([0,1,0])
count = len(b)
ls = list(set(b))
res = np.zeros((len(ls),a.shape[1]))
for i in ls:
    res[i] = np.array([a[x] for x in range(0,count) if b[x] == i]).sum(axis=0)
print res

我获得的打印结果如下：

[[ 5.  7.  9.]
 [ 2.  3.  4.]]

我想要做的是，由于b的第一个和第三个元素为0，所以我执行a[0]+a[2]，结果是[5, 7, 9]作为结果的一行。同样地，b的第二个元素是1，因此我执行a[1]，结果是[2, 3, 4]作为结果的另一行。

但是对于大型数组，我的实现似乎相当慢。是否有更好的实现方法？

我知道numpy中有一个bincount函数。但它似乎只支持1d数组。

谢谢大家帮忙！

- pfc

---

我不太明白你想做什么。 - void

@s_vishnu，我已经更新了我的问题。谢谢。 - pfc

所以您想要添加所有具有相应“0”的行和所有具有“1”的行吗？ - void

@s_vishnu 是的。但更普遍的是，可能会有超过2个id。 - pfc

Answer 1

numpy_indexed包（免责声明：我是它的作者）旨在以高效矢量化和通用的方式解决此类问题：

import numpy_indexed as npi
unique_b, mean_a = npi.group_by(b).mean(a)

请注意，此解决方案是通用的，因为它提供了丰富的标准缩减函数（sum，min，mean，median，argmin等），如果您需要使用不同的轴，则提供axis关键字。此外，还可以按比仅限于正整数数组的更复杂的方式进行分组，例如任意dtype的多维数组的元素。

import numpy_indexed as npi
# this caches the complicated O(NlogN) part of the operations
groups = npi.group_by(b)
# all these subsequent operations have the same low vectorized O(N) cost
unique_b, mean_a = groups.mean(a)
unique_b, sum_a = groups.sum(a)
unique_b, min_a = groups.min(a)

- Eelco Hoogendoorn

---

还有非常感谢你。我已经测试过了。你的实现是我测试过的所有方法中最快的。谢谢你，我强烈推荐这个实现！！！ - pfc

很高兴能够提供帮助。这些函数在内部使用reduceat进行缩减，对于我的基准测试来说，它比addat快得多（这可能是一些numpy实现细节；我没有看到其根本原因）。因此，Divakar的更新的reduceat解决方案应该具有类似的性能水平，因为它经过与此库相同的步骤。 - Eelco Hoogendoorn

Answer 2

方法一

您可以使用np.add.at，它适用于通用维度的ndarrays，而不像np.bincount仅适用于1D数组 -

np.add.at(res, b, a)

样例运行 -

In [40]: a
Out[40]: 
array([[1, 2, 3],
       [2, 3, 4],
       [4, 5, 6]])

In [41]: b
Out[41]: array([0, 1, 0])

In [45]: res = np.zeros((b.max()+1, a.shape[1]), dtype=a.dtype)

In [46]: np.add.at(res, b, a)

In [47]: res
Out[47]: 
array([[5, 7, 9],
       [2, 3, 4]])

为了计算平均值，我们需要使用np.bincount获取每个标签/标记的计数，然后沿每行除以它们，如下所示 -

In [49]: res/np.bincount(b)[:,None].astype(float)
Out[49]: 
array([[ 2.5,  3.5,  4.5],
       [ 2. ,  3. ,  4. ]])

如果我们希望处理的 b 不一定是从 0 开始的连续序列，我们可以将其泛化并编写一个漂亮的小函数来更加清晰地处理求和和平均值，如下所示 -

def groupby_addat(a, b, out="sum"):
    unqb, tags, counts = np.unique(b, return_inverse=1, return_counts=1)
    res = np.zeros((tags.max()+1, a.shape[1]), dtype=a.dtype)
    np.add.at(res, tags, a)

    if out=="mean":
        return unqb, res/counts[:,None].astype(float)
    elif out=="sum":
        return unqb, res
    else:
        print "Invalid output"
        return None

示例运行 -

In [201]: a
Out[201]: 
array([[1, 2, 3],
       [2, 3, 4],
       [4, 5, 6]])

In [202]: b
Out[202]: array([ 5, 10,  5])

In [204]: b_ids, means = groupby_addat(a, b, out="mean")

In [205]: b_ids
Out[205]: array([ 5, 10])

In [206]: means
Out[206]: 
array([[ 2.5,  3.5,  4.5],
       [ 2. ,  3. ,  4. ]])

方法 #2

我们也可以使用np.add.reduceat，这可能会更加高效 -

def groupby_addreduceat(a, b, out="sum"):
    sidx = b.argsort()
    sb = b[sidx]
    spt_idx =np.concatenate(([0], np.flatnonzero(sb[1:] != sb[:-1])+1, [sb.size]))
    sums = np.add.reduceat(a[sidx],spt_idx[:-1])

    if out=="mean":
        counts = spt_idx[1:] - spt_idx[:-1]
        return sb[spt_idx[:-1]], sums/counts[:,None].astype(float)
    elif out=="sum":
        return sb[spt_idx[:-1]], sums
    else:
        print "Invalid output"
        return None

样例运行 -

In [201]: a
Out[201]: 
array([[1, 2, 3],
       [2, 3, 4],
       [4, 5, 6]])

In [202]: b
Out[202]: array([ 5, 10,  5])

In [207]: b_ids, means = groupby_addreduceat(a, b, out="mean")

In [208]: b_ids
Out[208]: array([ 5, 10])

In [209]: means
Out[209]: 
array([[ 2.5,  3.5,  4.5],
       [ 2. ,  3. ,  4. ]])

- Divakar

---

是的，就是这样！顺便说一下，我还需要以相同的方式计算平均值。是否有类似的平均值函数？谢谢！ - pfc

@pfc 请查看修改内容。我已经自行编辑了问题标题。希望这样可以。如有需要，请随意进一步编辑。 - Divakar

谢谢。我又遇到了一个问题。实际上，数组 b 可能不是 [0, 1]，而是 [100, 101]。在这种情况下，我遇到了越界错误。我该怎么办？非常感谢您的耐心帮助。 - pfc

@pfc 你是否按照这篇帖子中的建议进行了初始化 - res = np.zeros((b.max()+1, a.shape[1]), dtype=a.dtype)？ - Divakar

谢谢。你解决了我的问题。感谢你的善良！ - pfc