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

Admin 1
Viewed 339

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.

list python
3 Answers
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)]
user_29473780

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

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)
user_123793250
Related