How to obtain only the first True value from each row of a numpy array?
You can use cumsum
, and find the first bool by comparing the result with 1.
all_bools.cumsum(axis=1).cumsum(axis=1) == 1
array([[False, True, False],
[ True, False, False],
[False, False, True],
[False, False, False]])
This also accounts for the issue @a_guest pointed out. The second cumsum
call is needed to avoid matching all False
values between the first and second True
value.
If performance is important, use argmax
and set values:
y = np.zeros_like(all_bools, dtype=bool)
idx = np.arange(len(x)), x.argmax(axis=1)
y[idx] = x[idx]
y
array([[False, True, False],
[ True, False, False],
[False, False, True],
[False, False, False]])
Perfplot Performance Timings
I'll take this opportunity to show off perfplot
, with some timings, since it is good to see how our solutions vary with different sized inputs.
import numpy as np
import perfplot
def cs1(x):
return x.cumsum(axis=1).cumsum(axis=1) == 1
def cs2(x):
y = np.zeros_like(x, dtype=bool)
idx = np.arange(len(x)), x.argmax(axis=1)
y[idx] = x[idx]
return y
def a_guest(x):
b = np.zeros_like(x, dtype=bool)
i = np.argmax(x, axis=1)
b[np.arange(i.size), i] = np.logical_or.reduce(x, axis=1)
return b
perfplot.show(
setup=lambda n: np.random.randint(0, 2, size=(n, n)).astype(bool),
kernels=[cs1, cs2, a_guest],
labels=['cs1', 'cs2', 'a_guest'],
n_range=[2**k for k in range(1, 8)],
xlabel='N'
)
The trend carries forward to larger N. cumsum
is very expensive, while there is a constant time difference between my second solution, and @a_guest's.
You can use the following approach using np.argmax
and a product with np.logical_or.reduce
for dealing with rows that are all False
:
b = np.zeros_like(a, dtype=bool)
i = np.argmax(a, axis=1)
b[np.arange(i.size), i] = np.logical_or.reduce(a, axis=1)
Timing results
Different versions in increasing performance, i.e. fastest approach comes last:
In [1]: import numpy as np
In [2]: def f(a):
...: return a.cumsum(axis=1).cumsum(axis=1) == 1
...:
...:
In [3]: def g(a):
...: b = np.zeros_like(a, dtype=bool)
...: i = np.argmax(a, axis=1)
...: b[np.arange(i.size), i] = np.logical_or.reduce(a, axis=1)
...: return b
...:
...:
In [4]: x = np.random.randint(0, 2, size=(1000, 1000)).astype(bool)
In [5]: %timeit f(x)
10.4 ms ± 155 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [6]: %timeit g(x)
120 µs ± 184 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [7]: def h(a):
...: y = np.zeros_like(x)
...: idx = np.arange(len(x)), x.argmax(axis=1)
...: y[idx] += x[idx]
...: return y
...:
...:
In [8]: %timeit h(x)
92.1 µs ± 3.51 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [9]: def h2(a):
...: y = np.zeros_like(x)
...: idx = np.arange(len(x)), x.argmax(axis=1)
...: y[idx] = x[idx]
...: return y
...:
...:
In [10]: %timeit h2(x)
78.5 µs ± 353 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)