Question

Python的numpy库对花式索引进行“矢量化”有时会产生意想不到的结果。例如：

import numpy
a = numpy.zeros((1000,4), dtype='uint32')
b = numpy.zeros((1000,4), dtype='uint32')
i = numpy.random.random_integers(0,999,1000)
j = numpy.random.random_integers(0,3,1000)

a[i,j] += 1
for k in xrange(1000):
    b[i[k],j[k]] += 1

在数组'a'和'b'中给出不同的结果（即，无论重复如何，元组（i，j）的外观在'a'中显示为1，而重复在'b'中计数）。这很容易验证如下：

numpy.sum(a)
883
numpy.sum(b)
1000

值得注意的是，花式索引版本比for循环快了近两个数量级。我的问题是：“在提供的示例中使用for循环实现的numpy计算重复计数是否有效？”

Answer 1

这应该做你想要的：

np.bincount(np.ravel_multi_index((i, j), (1000, 4)), minlength=4000).reshape(1000, 4)

作为细分，ravel_multi_index将i和j指定的索引对转换为整数索引，转换为C-flattened数组; bincount计算每个值0..4000在索引列表中出现的次数;并且reshape将C-flattened数组转换回2d数组。

在性能方面，我的测量速度比“b”快200倍，比“a”快5倍;你的里程可能会有所不同。

由于您需要将计数写入现有数组a，请尝试以下操作：

u, inv = np.unique(np.ravel_multi_index((i, j), (1000, 4)), return_inverse=True)
a.flat[u] += np.bincount(inv)

我使第二种方法比“a”慢一点（2x），这并不太令人惊讶，因为unique阶段会变慢。

Python Numpy库中的复合赋值运算符

1 个答案: