Python, how to extract the column index of the same elements in an array and store them in new arrays with the same numbers inside

Question

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

Answer 1

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

Answer 2

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

Answer 3

You can use numpy.unique with the return_inverse and return_counts flags like so:

import numpy as np

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

unq,inv,cnt=np.unique(x,0,1,1)

dict(zip(unq,np.split(inv.argsort(),cnt[:-1].cumsum())))
# {0: array([ 9, 11]), 1: array([4, 5]), 2: array([10]), 3: array([ 0,  3, 12]), 5: array([6, 8]), 6: array([1]), 8: array([2, 7])}