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

【问题标题】：How to get a list of indexes selected by a specific value efficiently with numpy arrays?如何使用 numpy 数组有效地获取由特定值选择的索引列表？
【发布时间】：2018-02-08 13:15:06
【问题描述】：

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

# [...]

我想到的第一种方法（这不是很 Python）：

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 数组实现这一目标的最有效方法是什么？

【问题讨论】：

虽然这不是 numpy，但结合 itertools.groupby 和 enumerate 会起作用。
arr 总是排序吗？
@hpaulj 不，你可以看到数组中间还有一个 9。我添加了更多数字以避免误解
这个问题本质上没有像数组一样的东西 - 它实际上是关于收集列表中的位置。您的结果是不同长度的列表，这是一个很好的指标，表明这没有快速的 numpy 解决方案。

标签： python python-3.x numpy

【解决方案1】：

可能不是禁食，但带有 numpy 的 oneliner 将是：

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

【讨论】：

【解决方案2】：

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

【讨论】：

collections.defaultdict 将最后一行缩短为 dd[val].append(i)。

【解决方案3】：

您可以使用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_* 列表。

【讨论】：