制作一个索引数组，其中包含某个值在另一个数组中的位置

Question

我试图创建一个Python列表，其中包含另一个整数列表（或Numpy数组）中元素的索引等于1的Python列表。我试图是这样的（对于1或2维的情况下）：

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

这也可以更紧凑地写为

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

但我想知道如果有一个甚至更紧凑的方式（或内置函数），可用于相同的目的。在纯Python，Numpy或其他地方有没有这种方便的功能？

Answer 1

可以使用numpy.where功能

在您提供的例子，其中的子表有不同的长度numpy.array检查这个帖子

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

是不是一种选择，因为你不能改变你的list到array（子列表长度应该相等）。这里是一个多解：

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

其中ls对应于子列表和e到子列表的元素

Answer 2

使用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