0

I know numpy.where gives a tuple of the array coordinates where the condition applies. But what if I want an array? assume the following 2d array:

a=np.array([[1 1 1 1 0],
            [1 1 1 0 0],
            [1 0 0 0 0],
            [1 0 1 1 1],
            [1 0 0 1 0]])

Now what I want is only the first occurrence of zeros, but for every row, even if it doesn't exist. Something like indexOf() in Java. So the output look like:

array([-1,2,2,1,0])

I need to cut pieces of an ndarray and it would be much easier to reduce a dimension rather than having a tuple and try to regenerate the missing rows.

anishtain4
  • 2,342
  • 2
  • 17
  • 21

2 Answers2

5

Is this what you are looking for?

import numpy as np

a=np.array([[1, 1, 1, 1, 0],
            [1, 1, 1, 0, 0],
            [1, 0, 0, 0, 0],
            [1, 0, 1, 1, 1],
            [1, 0, 0, 1, 0]])

np.argmax(a==0, axis=0) - ~np.any(a==0, axis=0)

Output:

array([-1,  2,  2,  1,  0], dtype=int64)

The idea here is that np.argmax finds the index of the first matching element in each column (axis=0 for columns, which appears to be what you want in the output, but if you actually want rows, use axis=1). Because np.argmax returns 0 for columns that do not match at all, I subtract 1 from the result for each column that doesn't contain any 0.

PaSTE
  • 4,050
  • 18
  • 26
0

Here is a less crafty solution but arguably easier to undestand. First finds all matches and then creates an array with the first element of the matches and -1 if len == 0.

a=np.array([[1,1,1,1,0],
        [1,1,1,0,0],
        [1,0,0,0,0],
        [1,0,1,1,1],
        [1,0,0,1,0]])

matches = [np.where(np.array(i)==0)[0] for i in a.T]
np.array([i[0] if len(i) else -1 for i in matches]) # first occurence, else -1

array([-1,  2,  2,  1,  0])
Anton vBR
  • 18,287
  • 5
  • 40
  • 46