Python, how to extract the column index of the same elements in an array and store them in new array-CodePudding

Let's say I have a 1D array and I want to obtain all the column indexes of the same numbers in this array and store them in separate arrays.

For example, all the numbers for 3 are in index 0,3,12, for 1 in index 4,5 etc.

x = np.array([3,6,8,3,1,1,5,8,5,0,2,0,3])

So the output would be

a = np.array([0,3,12]) # Number 3
b = np.array([1]) # Number 6
c = np.array([2,7]) # Number 8
d = np.array([4,5]) # Number 1
e = np.array([9,11]) # Number 0
f = np.array([6,8]) # Number 5
g = np.array([10]) # Number 2

CodePudding user response：

np.argwhere(x==3)

Outputs:

array([[ 0],
       [ 3],
       [12]])

Edit: if you want an 1D Array just use the flatten() method on the output

Most functions to search for indexes in numpy starts with arg for the next time ;)

CodePudding user response：

def same_num_index(arr):
    output = {k:[] for k in set(arr)}
    for i,x in enumerate(arr):
        output[x].append(i)
    return output

x = np.array([3,6,8,3,1,1,5,8,5,0,2,0,3])

index_of = same_num_index(x)

index_of[3]

produces:

[0, 3, 12]

Edit:

@jack's answer uses numpy built-in function and should be preferred