numpy 获取特定数值的索引_如何使用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]

# [...]

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

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

解决方案

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)