Question

我有一个二维数组a和一个数组b。我想根据a中的每个id计算数组b组中的行总和。例如：

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的第1和第3个元素是0，我执行a[0]+a[2]，[5, 7, 9]作为结果的一行。同样，b的第二个元素是1，因此我执行a[1]，[2, 3, 4]作为结果的另一行。

但是对于大型阵列来说，我的实现似乎相当慢。有没有更好的实施？

我知道bincount中有numpy个功能。但似乎只支持1d阵列。谢谢大家的帮助！

Answer 1

numpy_indexed包（免责声明：我是它的作者）是为了以有效的矢量化和一般方式解决这类问题：

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

请注意，此解决方案是通用的，因为它提供了一组丰富的标准缩减函数（sum，min，mean，median，argmin等），轴关键字，如果您需要使用不同的轴，以及也可以通过比正整数数组更复杂的事物进行分组，例如任意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)

Answer 2

方法＃1

你可以使用np.add.at，它适用于通用维度的ndarray，不像只需要1D数组的np.bincount -

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]])

要计算mean值，我们需要使用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的顺序，我们可以使它成为通用的并放入一个很好的小函数来以更干净的方式处理求和和平均值，如下所示 - < / p>

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. ]])

如何通过id有效地求和和表示2D NumPy数组？

2 个答案: