35

I need to figure out how I can find all the index of a value in a 2d numpy array.

For example, I have the following 2d array:

([[1 1 0 0],
  [0 0 1 1],
  [0 0 0 0]])

I need to find the index of all the 1's and 0's.

1: [(0, 0), (0, 1), (1, 2), (1, 3)]
0: [(0, 2), (0, 3), (1, 0), (1, 1), (the entire all row)]

I tried this but it doesn't give me all the indexes:

t = [(index, row.index(1)) for index, row in enumerate(x) if 1 in row]

Basically, it gives me only one of the index in each row [(0, 0), (1, 2)].

Alex Riley
  • 169,130
  • 45
  • 262
  • 238
Pete
  • 754
  • 2
  • 13
  • 27

3 Answers3

69

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.

VIX
  • 605
  • 4
  • 15
Alex Riley
  • 169,130
  • 45
  • 262
  • 238
  • 4
    Just a note: `zip(*x)` can be slow for large `x`, where `x = np.where(a == 1)`. If a 2d array is fine, `np.asarray(x).T` is quite a bit faster, and if you *really* want a list-of-lists, `np.asarray(x).T.tolist()` is marginally faster. – jme Nov 27 '14 at 17:41
  • Is it possible to do something similar but with the index where the 2d array is nearest to a given value, e.g. np.where_is_nearest_to(0.99)? – falematte Jun 02 '23 at 17:17
19

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]]
Guillaume Lebreton
  • 2,586
  • 16
  • 25
15

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.

Mike
  • 6,813
  • 4
  • 29
  • 50