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

  • have a look at the updates in my responce and check one more solution for the case when sublists have unequal length – Luchko Apr 22 '17 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 :) –  Apr 22 '17 at 20:55

2 Answers2

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

Community
  • 1
  • 1
Luchko
  • 1,123
  • 7
  • 15
  • Thanks, it works :) (I initially had trouble with 2D case, but I needed to convert the input B array first to a numpy array...) –  Apr 22 '17 at 19:52
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
gold_cy
  • 13,648
  • 3
  • 23
  • 45
  • 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)) –  Apr 22 '17 at 20:00
  • you can also write ``zip(*id)`` – Luchko Apr 22 '17 at 20:08
  • I see... this makes the code very simple. thanks again :) –  Apr 22 '17 at 20:14