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

Question

我想创建一个Python列表，其中包含另一个整数列表(或Numpy数组)中等于1的元素的索引。我正在尝试这样做(适用于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或其他地方是否有这样方便的函数...

- user4834273

---

请查看我的回复中的更新，并检查一种更多的解决方案，以处理子列表长度不相等的情况。 - Luchko

@Luchko 谢谢你的更新，我稍后也会仔细研究它。我正在创建一个索引数组用于一些聚类分析，其大小是先验未知的，因此我使用了 Python 列表。但我也可以使用 Numpy 数组进行索引（或将列表转换为数组），所以没问题。谢谢 :) - user4834273

Answer 1

你可以使用numpy.where函数。

查看此帖子。

从NumPy数组中查找等于零的元素的索引

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对应于子列表中的元素

- Luchko

---

谢谢，它可以工作 :) （最初我在2D情况下遇到了问题，但我需要先将输入B数组转换为numpy数组...） - user4834273

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

- gold_cy

---

非常感谢！从获得的索引数组中提取索引的常见方法是像“for (i,j) in zip(id[0],id[1]): …”这样吗？（其中id = np.where(arr == 1)） - user4834273

你也可以写成 zip(*id)。 - Luchko

我明白了...这让代码变得非常简单。再次感谢 :) - user4834273