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

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

这也可以写得更简洁，如

#--- 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或其他地方有这样方便的函数吗？

您可以使用

Numpy.where

函数

查看此帖子

在您提供的示例中，子列表的长度不同。

numpy。数组

不是一个选项，因为您无法将

列表

转换为

数组

（子列表的长度应相等）。这里还有一个解决方案：

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

对应于子列表的元素

使用

numpy。其中

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

---

谢谢，它可以工作：）（我最初在2D case上遇到问题，但我需要先将输入B数组转换为numpy数组…）非常感谢！这是否可能是一种从获得的索引数组中提取索引的常用方法，如zip中的“for（i，j）（id[0]，id[1]）：…”？（其中id=np.where（arr==1））您还可以编写

zip（*id）

我明白了。。。这使得代码非常简单。再次感谢：）请查看“我的响应”中的更新，并在子列表不一致的情况下再检查一个解决方案length@Luchko谢谢你的更新，我以后也会仔细研究的。我制作的索引数组是用于一些聚类分析的，它的大小是未知的，所以我使用了Python列表。但是我也可以使用Numpy数组进行索引（或者转换列表->数组），所以没有问题。谢谢：）

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

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