製作一個索引數組，其中包含某個值在另一個數組中的位置

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