python - 有效地找到数组中所有值的索引

Question

我有一个非常大的数组，由 0 到 N 之间的整数组成，其中每个值至少出现一次。

我想知道，对于每个值 k，我数组中的所有索引，其中数组的值等于 k。

例如:

arr = np.array([0,1,2,3,2,1,0])
desired_output = {
    0: np.array([0,6]),
    1: np.array([1,5]),
    2: np.array([2,4]),
    3: np.array([3]),
    }

现在，我通过在 range(N+1) 上循环并调用 np.where N 次来完成此操作。

indices = {}
for value in range(max(arr)+1):
    indices[value] = np.where(arr == value)[0]

这个循环是迄今为止我的代码中最慢的部分。 (arr==value 评估和 np.where 调用都占用了大量时间。)是否有更有效的方法来做到这一点？

我也尝试过使用 np.unique(arr, return_index=True) 但这只会告诉我第一个索引，而不是所有索引。

Answer 1

方法 #1

这是一种将这些索引作为数组列表获取的矢量化方法 -

sidx = arr.argsort()
unq, cut_idx = np.unique(arr[sidx],return_index=True)
indices = np.split(sidx,cut_idx)[1:]

如果您想要将每个唯一元素与其索引相对应的最终字典，最后我们可以使用循环理解 -

dict_out = {unq[i]:iterID for i,iterID in enumerate(indices)}

---

方法 #2

如果您只对数组列表感兴趣，这里有一个用于提高性能的替代方案 -

sidx = arr.argsort()
indices = np.split(sidx,np.flatnonzero(np.diff(arr[sidx])>0)+1)