4

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...?

Improve this question
2
  • have a look at the updates in my responce and check one more solution for the case when sublists have unequal length Commented Apr 22, 2017 at 20:29
  • @Luchko Thanks for the update, I will study it also closely later. The index array I'm making is for some cluster analysis, whose size is unkonwn a priori so I used a Python list. But I could also use Numpy arrays for indexing (or convert lists -> arrays) so no problem. thanks :) Commented Apr 22, 2017 at 20:55

2 Answers 2

Reset to default
3

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

Improve this answer
Sign up to request clarification or add additional context in comments.

1 Comment

Thanks, it works :) (I initially had trouble with 2D case, but I needed to convert the input B array first to a numpy array...)
1

Use numpy.where

arr = np.arange(16).reshape(4,4)

>>array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])

np.where(arr == 1)

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

arr[0][1]

>>1
Improve this answer

3 Comments

Thanks very much! Is it possibly a common way to extract indices from the obtained index array like "for (i,j) in zip(id[0],id[1]): ..." ? (where id = np.where(arr == 1))
you can also write zip(*id)
I see... this makes the code very simple. thanks again :)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.