如何使用numpy数组有效地获取由特定值选择的索引列表? [英] How to get a list of indexes selected by a specific value efficiently with numpy arrays?

问题描述

我有一个像这样的numpy数组:

I have a numpy array like this:

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])

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

And I want to get a list of indexes of the found values by groups

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):

The first approach that comes to my mind (that´s not very 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数组哪种方法最有效?

Which is the most efficient way to achieve this with numpy arrays?

推荐答案

这应该不会太慢.该数组仅迭代一次. 结果(ind)是字典值->索引列表.

This should not be too slow. The array is iterated only once. The result (ind) is a dictionary value -> list of indexes.

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)

这篇关于如何使用numpy数组有效地获取由特定值选择的索引列表?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！