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

前端 未结 3 1358
庸人自扰
庸人自扰 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:30

    You can use numpy.unique to find all the unique values and numpy.where to find their indexes:

    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
    

    Now you have a dictionary of indexes of each value and you can assign what you want to your index_list_* lists.

提交回复
热议问题