Making an index array that contains the position of some value in another array

Question

I am trying to create a Python list that contains indices of the elements equal to 1 in another integer list (or Numpy array). What I am trying is something like this (for either 1- or 2-dimensional case):

#--- 1D case ---
A = [ 1, 0, 0, 1, 1 ]
idx = []
for i in range( len( A ) ):
    if A[ i ] == 1 : idx.append( i )

print( idx )   # [ 0, 3, 4 ]

#--- 2D case --- 
B = [ [ 1, 0, 0, 1, 1 ], [ 0, 1, 1 ] ]
idx2 = [ [] for i in range( len( B ) ) ]

for i in range( len( B ) ):
    for j in range( len( B[ i ] ) ):
        if B[ i ][ j ] == 1 : idx2[ i ].append( j )

print( idx2 )   #[ [0,3,4], [1,2] ]

This may also be written more compactly as

#--- 1D case ---
idx = [ i for i in range( len(A) ) if A[ i ] == 1 ]

#--- 2D case ---
idx2 = []
for i in range( len( B ) ):
    tmp = [ k for k in range( len(B[ i ]) ) if B[ i ][ k ] == 1 ]
    idx2.append( tmp )

But I am wondering if there is an even more compact way (or builtin function) that can be used for the same purpose. Is there any such convenient function in pure Python, Numpy, or elsewhere...?

Answer 1

you can use numpy.where function

check this post

Find indices of elements equal to zero from numpy array

import numpy as np

#-----1D case------
A = np.array([0,1,2,3,1,4,5,1,2])
print(np.where(A==1))

>>> (array([1, 4, 7]),)

#-----2D case------
A = np.array([[0,1,2,3],[1,2,3,5],[1,2,3,1]])
print(np.where(A==1))

>>> (array([0, 1, 2, 2]), array([1, 0, 0, 3]))

in the examples you provided where the sublists have different length numpy.array is not an option as you cannot transform your list to array (sublists length should be equal). here is one more solution:

B = [ [ 1, 0, 0, 1, 1 ], [ 0, 1, 1 ] ]
inds = [(i,j) for j,ls in enumerate(B) for i,e in enumerate(ls) if e==1]
print(inds)

>>>> [(0, 0), (3, 0), (4, 0), (1, 1), (2, 1)]

where ls corresponds to sublist and e to the element of the sublist