How to get a list of indexes selected by a specific value efficiently with numpy arrays?

前端 未结 3 1388
庸人自扰
庸人自扰 2021-01-15 12:47

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

3条回答
  •  独厮守ぢ
    2021-01-15 13:29

    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)
    

提交回复
热议问题