23

I have a NumPy array:

[[  0.   1.   2.   3.   4.]
 [  7.   8.   9.  10.   4.]
 [ 14.  15.  16.  17.   4.]
 [  1.  20.  21.  22.  23.]
 [ 27.  28.   1.  20.  29.]]

which I want to quickly find the coordinates of specific values and avoid Python loops on the array. For example the number 4 is on:

row 0 and col 4
row 1 and col 4
row 2 and col 4

and a search function should return a tuple:

((0,4),(1,4),(2,4))

Can this be done directly via NunmPy's functions?

Georgy
  • 12,464
  • 7
  • 65
  • 73
Ηλίας
  • 2,560
  • 4
  • 30
  • 44

2 Answers2

30

If a is your array, then you could use:

ii = np.nonzero(a == 4)

or 

ii = np.where(a == 4)

If you really want a tuple, you can convert from the tuple of arrays to the tuple of tuples, but the return value from the numpy functions is convient for then doing other operations on your array.

Conversion to a tuple for the OP's specification:

tuple(zip(*ii))
JoshAdel
  • 66,734
  • 27
  • 141
  • 140
24
a = numpy.array([[  0.,  1.,  2.,  3.,  4.],
                 [  7.,  8.,  9., 10.,  4.],
                 [ 14., 15., 16., 17.,  4.],
                 [  1., 20., 21., 22., 23.],
                 [ 27., 28.,  1., 20., 29.]])
print numpy.argwhere(a == 4.)

prints

[[0 4]
 [1 4]
 [2 4]]

The usual caveats for floating point comparisons apply.

Sven Marnach
  • 574,206
  • 118
  • 941
  • 841
  • 1
    Just as a note from the documentation "The output of `np.argwhere` is not suitable for indexing arrays. For this purpose use `np.where(a)` instead." – JoshAdel Mar 30 '11 at 14:17
  • @JoshAdel: Why do you think the OP wants to use this for indexing? I chose the function that produces output that is closest to the OP's desired output. – Sven Marnach Mar 30 '11 at 14:25
  • @SvenMarnach: Could have gone with both, he just mentioned first the `.where` – Ηλίας Mar 30 '11 at 14:28
  • @SvenMarnach: Don't be sad, you still have 17k rep points on me :-) But seriously, I also provided a means for getting the OP's desired output, but I also wanted to point out something about how numpy works that might be useful for others. It wasn't meant to be a knock on your solution, but some additional (and what I thought was useful) info. – JoshAdel Mar 30 '11 at 14:35
  • @SvenMarnach, @JoshAdel It was also part concurrency issue, I think SO would have benefited if Google Wave was successful! – Ηλίας Mar 30 '11 at 14:45
  • @JoshAdel: I think your comment was helpful. My point only was that I try not to make assumptions on the OP's goals if not necessary. – Sven Marnach Mar 30 '11 at 14:45