Numpy get coordinates where condition is satisfied together with coordinates?

Question

Using numpy, what's the fastest way to get the coordinates of elements of a 2D array that satisfy some condition both in terms of the cell value and in terms of the coordinates?

For example, say I have this array:

[[45 78 95 49 98 18 54]
 [36 77 92 15 68 25 89]
 [29 44 58 18 85 65 43]
 [19 63 62 78 48 12 43]
 [90 65 17 90 65 44 53]]

and say I want to get the coordinates of all elements > 50 that are in columns 1 and 3 and eventrows. I could do something like:

• Select all columns but 1, 3 and all rows except even rows

• Zero them out

  [[ 0 78  0 49  0  0  0]
   [36 77 92 15 68 25 89]
   [ 0 44  0 18  0  0  0]
   [19 63 62 78 48 12 43]
   [ 0 65  0 90  0  0  0]]
  

• Get the coordinates of all elements from the resulting array that are > 50

   [[0 1]
    [1 1]
    [1 2]
    [1 4]
    [1 6]
    [3 1]
    [3 2]
    [3 3]
    [4 1]
    [4 3]]
  

with something like this:

rs = numpy.array([0, 2, 4])[:, None]
cs = numpy.array([0, 2, 4, 5, 6])
a[rs, cs] = 0
res = numpy.argwhere(a > 50)

Is there a faster (since I'd have to copy a before zeroing out, since I need its values later) / shorter (esp. in terms of being more numpy-y) way?

Answer 1

Here is one method using ogrid:

>>> i, j = np.ogrid[(*map(slice, a.shape),)]
>>> np.argwhere((a>50) & ((i|2==3) | (j|2==3)))
array([[0, 1],
       [1, 1],
       [1, 2],
       [1, 4],
       [1, 6],
       [3, 1],
       [3, 2],
       [3, 3],
       [4, 1],
       [4, 3]])