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

Question

我有一个像这样的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]

# [...]

我想到的第一种方法（不是非常pythonic）：

i = 0
for x in arr:
    if x == 2:
        index_list_2 += [i]
    if x == 3:
        index_list_3 += [i]
    if x == 4:
        index_list_4 += [i]
    if x == 9:
        index_list_9 += [i]
    i += 1

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

Answer 1

这不应该太慢。该数组只迭代一次。结果（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)

Answer 2

您可以使用numpy.unique查找所有唯一值，并使用numpy.where查找其索引：

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

现在您有一个每个值的索引字典，您可以将您想要的内容分配给index_list_*列表。

Answer 3

可能不是最快但是一个有numpy的衬里将是：

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