基于Numpy，根据另一个数组对数组中的数据进行求和

Question

我有两个2D的numpy数组（在这个例子中简化了大小和内容），它们的大小相同。

一个ID矩阵：

1 1 1 2 2
1 1 2 2 5
1 1 2 5 5
1 2 2 5 5
2 2 5 5 5

一个值矩阵：

14.8 17.0 74.3 40.3 90.2
25.2 75.9  5.6 40.0 33.7
78.9 39.3 11.3 63.6 56.7
11.4 75.7 78.4 88.7 58.6
79.6 32.3 35.3 52.5 13.3

我的目标是按照第一个矩阵中的ID分组，对第二个矩阵中的值进行计数和求和。

1: (8, 336.8)
2: (9, 453.4)
5: (8, 402.4)

我可以用 for 循环来完成这个任务，但当矩阵的大小达到数千而不仅仅是 5x5，并且有数千个唯一的 ID 时，处理起来需要很长时间。

numpy 是否有聪明的方法或组合方法来完成这个任务？

- Chau

Answer 1

这是一种向量化的方法，通过使用 np.unique 和 np.bincount 组合来获取 ID 以及基于 ID 的 value 总和值的计数。

unqID,idx,IDsums = np.unique(ID,return_counts=True,return_inverse=True)

value_sums = np.bincount(idx,value.ravel())

要将最终输出作为字典获取，您可以使用循环推导来收集总和值，如下所示 -

{i:(IDsums[itr],value_sums[itr]) for itr,i in enumerate(unqID)}

样例运行 -

In [86]: ID
Out[86]: 
array([[1, 1, 1, 2, 2],
       [1, 1, 2, 2, 5],
       [1, 1, 2, 5, 5],
       [1, 2, 2, 5, 5],
       [2, 2, 5, 5, 5]])

In [87]: value
Out[87]: 
array([[ 14.8,  17. ,  74.3,  40.3,  90.2],
       [ 25.2,  75.9,   5.6,  40. ,  33.7],
       [ 78.9,  39.3,  11.3,  63.6,  56.7],
       [ 11.4,  75.7,  78.4,  88.7,  58.6],
       [ 79.6,  32.3,  35.3,  52.5,  13.3]])

In [88]: unqID,idx,IDsums = np.unique(ID,return_counts=True,return_inverse=True)
    ...: value_sums = np.bincount(idx,value.ravel())
    ...: 

In [89]: {i:(IDsums[itr],value_sums[itr]) for itr,i in enumerate(unqID)}
Out[89]: 
{1: (8, 336.80000000000001),
 2: (9, 453.40000000000003),
 5: (8, 402.40000000000003)}

- Divakar

---

不错！我之前不知道np.unique函数有return_*参数。 - MB-F

@Divakar：谢谢！这正是我寻找的解决方案，由于矢量化而具有良好的性能。 - Chau

Answer 2

这可以通过以下几种简单方法的组合实现：

1. 使用 numpy.unique 查找每个 ID
2. 为每个 ID 创建一个布尔掩码
3. 对掩码中的 1 进行求和（计数）以及掩码为 1 的值

这可能看起来像这样：

import numpy as np

ids = np.array([[1, 1, 1, 2, 2],
                [1, 1, 2, 2, 5],
                [1, 1, 2, 5, 5],
                [1, 2, 2, 5, 5],
                [2, 2, 5, 5, 5]])

values = np.array([[14.8, 17.0, 74.3, 40.3, 90.2],
                   [25.2, 75.9,  5.6, 40.0, 33.7],
                   [78.9, 39.3, 11.3, 63.6, 56.7],
                   [11.4, 75.7, 78.4, 88.7, 58.6],
                   [79.6, 32.3, 35.3, 52.5, 13.3]])


for i in np.unique(ids):  # loop through all IDs
    mask = ids == i  # find entries that match current ID
    count = np.sum(mask)  # number of matches
    total = np.sum(values[mask])  # values of matches
    print('{}: ({}, {:.1f})'.format(i, count, total))  #print result

# Output:
# 1: (8, 336.8)
# 2: (9, 453.4)
# 5: (8, 402.4)

- MB-F

---

我指的是那个讨厌的for循环，我应该在问题中更清楚地说明。 - Chau

我认为没有一种简洁的方法可以在不使用for循环的情况下完成这个任务。虽然可能有其他方式，但很可能会导致代码非常难以阅读。如果你只有几个唯一的ID，使用for循环应该不会对性能造成太大影响。无论如何，我会再考虑一段时间... - MB-F

看起来我在Divakar的回答中被证明是错误的。 - MB-F

Answer 3

numpy_indexed 包（免责声明：我是它的作者）具有以优雅和向量化的方式解决这些问题的功能：

import numpy_indexed as npi
group_by = npi.group_by(ID.flatten())
ID_unique, value_sums = group_by.sum(value.flatten())
ID_count = groupy_by.count    

注意：如果您想计算总和和计数以便计算平均值，还可以使用group_by.mean；此外还有许多其他有用的功能。

- Eelco Hoogendoorn