Création d'un tableau d'index contenant la position d'une valeur dans un autre tableau

Question

J'essaie de créer une liste Python contenant les indices des éléments égaux à 1 dans une autre liste d'entiers (ou tableau Numpy). Ce que je suis en train est quelque chose comme ça (soit pour 1 ou 2 cas dimensions):

#--- 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] ] 

Cela peut également être écrit plus compacte

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

Mais je me demande s'il y a une manière encore plus compacte (ou fonction intégrée) qui peut être utilisée dans le même but. Y at-il une telle fonction pratique en Python pur, Numpy, ou ailleurs ...?

Answer 1

vous pouvez utiliser numpy.where fonction

vérifier ce poste

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

dans les exemples que vous avez fournis où les sous-listes ont une longueur différente numpy.array est pas une option que vous ne pouvez pas transformer votre list à array (la longueur des sous-listes doit être égale). voici une autre 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)] 

où ls correspond à subList et e à l'élément de la sous-liste

Answer 2

Utilisez 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