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

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

这也可以写得更简洁

#--- 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函数

查看这篇文章

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（子列表长度应该相等）。这是另一种解决方案：

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