如何從二維列表創建索引列表？

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)