CODEWITHSUNDEEP
pythonnumpy
Ηλίας
Ηλίας

Reputation: 2640

Coordinates of item on NumPy array

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?

Upvotes: 24

Views: 45292

Answers (2)

JoshAdel
JoshAdel

Reputation: 68682

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

Upvotes: 31

Sven Marnach
Sven Marnach

Reputation: 601401

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.

Upvotes: 25

Related Questions