How to create a list of indices from a 2D list?-CodePudding

If I have a 2 dimensional list, how can I generate a list of indices for a specific element in that 2D list?

For example, if I have the list

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

How could I make a list with this format

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

which is the list of 2 dimensional indexes of all 1s in two_d_list. The format of

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

would also work. I just need to retrieve the indices.

CodePudding user response：

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)

Result:

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

CodePudding user response：

With list comprehension:

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

CodePudding user response：

Working with list of lists can be cumbersome and quite slow. If you are able to use numpy arrays in your application, the solution becomes quite easy and fast.

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)