How to find the index of a value in 2d array in Python?
The problem with the list comprehension you provided is that it only goes one level deep, you need a nested list comprehension:
a = [[1,0,1],[0,0,1], [1,1,0]]
>>> [(ix,iy) for ix, row in enumerate(a) for iy, i in enumerate(row) if i == 0]
[(0, 1), (1, 0), (1, 1), (2, 2)]
That being said, if you are working with a numpy array, it's better to use the built in functions as suggested by ajcr.
Using numpy, argwhere may be the best solution:
import numpy as np
array = np.array([[1, 1, 0, 0],
[0, 0, 1, 1],
[0, 0, 0, 0]])
solutions = np.argwhere(array == 1)
print(solutions)
>>>
[[0 0]
[0 1]
[1 2]
[1 3]]
You can use np.where
to return a tuple of arrays of x and y indices where a given condition holds in an array.
If a
is the name of your array:
>>> np.where(a == 1)
(array([0, 0, 1, 1]), array([0, 1, 2, 3]))
If you want a list of (x, y) pairs, you could zip
the two arrays:
>>> list(zip(*np.where(a == 1)))
[(0, 0), (0, 1), (1, 2), (1, 3)]
Or, even better, @jme points out that np.asarray(x).T
can be a more efficient way to generate the pairs.