如何从二维列表创建索引列表？

Question

如果我有一个二维列表，如何为该二维列表中的特定元素生成索引列表？

例如，如果我有列表

two_d_list = [[0, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 1]]

我怎么能用这种格式制作一个列表

index_list = [[0, 1], [1, 0], [1, 1], [2, 3]]

这是two_d_list中所有 1 的二维索引列表。 的格式

index_list = [(0, 1), (1, 0), (1, 1), (2, 3)]

也会工作。 我只需要检索索引。

Answer 1

two_d_list = [[0, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 1]]

result = []
for i in range(len(two_d_list)):
    for j in range(len(two_d_list[i])):
        if two_d_list[i][j] == 1:
            result.append((i, j))

print(result)

结果：

[(0, 1), (1, 0), (1, 1), (2, 3)]

Answer 2

使用列表理解：

>>> [(r, c) for r, line in enumerate(two_d_list) for c, num in enumerate(line) if num==1]

[(0, 1), (1, 0), (1, 1), (2, 3)]

Answer 3

使用列表列表可能很麻烦而且很慢。 如果您能够在您的应用程序中使用numpy数组，那么该解决方案将变得非常简单和快速。

import numpy as np
two_d_list = [[0, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 1]]
# Create a numpy array
arr = np.array(two_d_list)
    
# np.where implements exactly the functionality you are after 
# then you stack the two lists of indices and transpose
indices = np.stack(np.where(arr == 1)).T
print(indices)