创建包含另一个数组中某个值位置的索引数组

Question

我试图创建一个Python列表，其中包含另一个整数列表（或Numpy数组）中等于1的元素的索引。我要做的是这样的（无论是一维还是二维）：

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

这也可以更简洁地写为

^{pr2}$

但我想知道是否有更紧凑的方法（或内置函数）可以用于同样的目的。在纯Python、Numpy或其他地方有这样方便的函数吗。。。？在

Answer 1

您可以使用numpy.where函数

检查这个帖子

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

在您提供的子列表具有不同长度的示例中，numpy.array不是一个选项，因为您无法将list转换为array（子列表长度应该相等）。还有一个解决方案：

^{pr2}$

其中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