如何以矢量化的方式修改任意索引的numpy数组?

Question

简化故事

假设我有一个数组arr和下标idx。对于idx中出现的每一个i，我希望arr[i]增加一个。

非矢量化的方法是这样的:

import numpy as np
arr = np.zeros(5)
idx = [0, 1, 1, 2, 0]
for i in idx:
arr[i] += 1

是否有办法将其矢量化?

注意arr[idx] += 1是无效的，因为索引重复。

arr = np.zeros(1)
idx = [0, 0]
arr[idx] += 1  # arr becomes array([1]), not array([2])

当然，在这个1D数组示例中使用np.unique()也可以达到同样的目的。但实际上我正在尝试处理二维数组，我怀疑计数元素将是最好的解决方案。

编辑

np.unique确实工作，但似乎有不必要的减速。我想要一个更快的方法(如果存在)。

这是一个2D索引的例子，10000个点没有重复。

arr = np.zeros((10000, 10000))
idx = np.stack([np.arange(10000), np.arange(10000)])
%timeit np.unique(idx, axis=1, return_counts=True)  # takes 1.93 ms
%timeit arr[idx[0], idx[1]] += 1  # takes 235 μs

显然，通过索引迭代要快10倍。

Edit2

@PaulS的答案比np.unique快。

%timeit np.add.at(arr, (idx[0], idx[1]), 1) # takes 925 μs

Edit3

下面是使用随机索引测试重复索引的示例。

arr = np.zeros((10000, 10000))
ran = (np.random.rand(10000)*10).astype(int)
idx = np.stack([ran, ran])
%timeit np.unique(idx, axis=1, return_counts=True)  # takes 3.24 ms
%timeit np.add.at(arr, (idx[0], idx[1]), 1) # takes 859 μs

(编辑:错字)<标题>

详细故事我正在尝试使用NumPy实现霍夫线变换算法。(我不使用cv2.HoughLines()的原因是因为我想直接从点的坐标得到结果，而不是从二进制数组)。

在(r, θ)平面上获得曲线很容易，但我在以矢量化的方式实现累加器时遇到了麻烦。目前我依靠的是将2D数据扁平化为1D。有没有更好更快的方法来进行积累?

提前感谢!

Answer 1

1D数组

另一个可能的解决方案:

np.add.at(arr, idx, 1)

输出:

[2. 2. 1. 0. 0.]

---

2 d数组

(谢谢，@mozway，对于你的例子，我现在在这里使用)

arr = np.zeros([3, 4], dtype=int)
idx = [[0, 0, 2, 0],
[1, 1, 3, 1]]
np.add.at(arr, (idx[0], idx[1]), 1)

输出:

array([[0, 3, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 1]])

Answer 2

使用numpy.unique获取唯一索引及其计数:

idx2, cnt = np.unique(idx, return_counts=True)
arr[idx2] += cnt

更新arr:

array([2, 2, 1, 0, 0])

with nd-arrays(以2D为例):

arr = np.zeros([3, 4], dtype=int)
idx = [[0, 0, 2, 0],
[1, 1, 3, 1]]
idx2, cnt = np.unique(idx, axis=1, return_counts=True)
arr[*idx2] = cnt

输出:

array([[0, 3, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 1]])

如果索引调换:

arr = np.zeros([3, 4], dtype=int)
idx = idx = [[0, 1], [0, 1], [2, 3], [0, 1]]
idx2, cnt = np.unique(idx, axis=0, return_counts=True)
arr[*idx2.T] = cnt