Python 如何使用numpy数组高效地获取由特定值选择的索引列表？

我有一个numpy数组，如下所示：

import numpy as np
arr = np.array([9, 6, 3, 8, 2, 3, 3, 4, 4, 9, 5, 6, 6, 6, 6, 7, 8, 9])

我想按组获得找到的值的索引列表

index_list_2 = [4 ]         # index list of the element with the value 2
index_list_3 = [2, 5, 6 ]
index_list_4 = [7, 8 ]
index_list_9 = [0, 9, 17]

# [...]

我想到的第一种方法（不是很像蟒蛇）：

使用numpy数组实现这一点的最有效方法是什么？

您可以使用以下方法查找所有唯一值并查找其索引：

import numpy as np
arr = np.array([2, 3, 3, 4, 4, 9, 5, 6, 6, 6, 6, 7, 8, 9])

# get the unique values
unique_arr = np.unique(arr)

# loop through the unique numbers and find the indeces
indexes_value = {}
for num in unique_arr:
    indexes = np.where(arr == num)[0]
    indexes_value[num] = indexes  # or list(indexes) if you prefer

---

现在您有了每个值的索引字典，您可以将您想要的内容分配给

索引列表*

列表。

这应该不会太慢。数组只迭代一次。 结果（ind）是一个字典值->索引列表

import numpy as np
arr = np.array([2, 3, 3, 4, 4, 9, 5, 6, 6, 6, 6, 7, 8, 9])

ind = dict()
for i, val in enumerate(arr):
  ind.setdefault(val, []).append(i)

---

可能不是fastes，但与numpy的一条航线是：

index_dict = {v: np.flatnonzero(arr == v) for v in np.unique(arr)}

---

虽然这不是麻木的，但是结合<代码>迭代器。GROPBY 和<代码>枚举会起作用。<代码> ARR 总是排序吗？@ HPAULJ不，正如您所看到的，在数组的中间也有9个。为了避免误解，我添加了更多的数字。这个问题没有什么本质上类似于数组的，它实际上是在列表中收集位置。您的结果是不同长度的列表，这很好地表明这没有快速的numpy解决方案。

collections.defaultdict

将最后一行缩短为

dd[val]。append（i）

。

index_dict = {v: np.flatnonzero(arr == v) for v in np.unique(arr)}