如何处理numpy.where条件不满意的情况？

Question

I am converting this array:

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

to: [2, 0, 1, 0, 0] .

Basically, I want to return the index of the first 1 in each sub-array. However, my problem is that I don't know how to handle the scenario where there is no 1 . I want it to return 0 if 1 is not found (like in my example).

The code below works fine but throws IndexError: index 0 is out of bounds for axis 0 with size 0 for the scenario I mentioned:

np.array([np.where(r == 1)[0][0] for r in x])

What is an easy way to handle this? It does not need to be restricted to numpy.where.

I am using Python 3 by the way.

Answer 1

Use mask of 1s and then argmax along each row to get the first matching index alongwith any to check for valid rows (rows with at least one 1 ) -

mask = x==1
idx = np.where(mask.any(1), mask.argmax(1),0)

Now, argmax on all False would return 0 . So, that plays right into the hands of the stated problem. As such, we can simply use mask.argmax(1) result. But in a general case, where the invalid specifier, let's call it invalid_val is not 0 , we can specify there inside np.where , like so -

idx = np.where(mask.any(1), mask.argmax(1),invalid_val)

Another method would be to get the first matching index on the mask and then index into the mask to see if any of the indexed values is False and set those as 0s -

idx = mask.argmax(1)
idx[~mask[np.arange(len(idx)), idx]] = 0 # or invalid_val

Answer 2

A simple modification to your code would be to add a condition to the list comprehension:

np.array([np.where(r == 1)[0][0] if 1 in r else 0 for r in x])
# 23.1 µs ± 43.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

A more concise and substantially faster way of obtaining the same result is:

np.argmax(x == 1, axis=1)
# 4.04 µs ± 45.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

or, equivalently:

np.argmin(x != 1, axis=1)
# 4.03 µs ± 13.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)