如何从 numpy 数组的每一行中仅获取第一个 True 值？

Question

我有一个 4x3 布尔 numpy 数组，我试图返回一个相同大小的数组，除了原始每一行上第一个 True 值的位置之外，它都是 False。所以如果我有一个

的起始数组

all_bools = np.array([[False, True, True],[True, True, True],[False, False, True],[False,False,False]])
all_bools
array([[False,  True,  True], # First true value = index 1
       [ True,  True,  True], # First true value = index 0
       [False, False,  True], # First true value = index 2
       [False, False, False]]) # No True Values

那我想回来

[[False, True, False],
 [True, False, False],
 [False, False, True],
 [False, False, False]]

所以前三行的索引 1、0 和 2 已设置为 True，仅此而已。基本上，原始方式中的任何 True 值（每行的第一个值之外）都已设置为 False。

我一直在用 np.where 和 np.argmax 来解决这个问题，但我还没有找到一个好的解决方案 - 任何帮助都非常感激。这需要运行很多次，所以我想避免迭代。

Answer 1

您可以使用cumsum，通过将结果与1进行比较来找到第一个布尔值。

all_bools.cumsum(axis=1).cumsum(axis=1) == 1 
array([[False,  True, False],
       [ True, False, False],
       [False, False,  True],
       [False, False, False]])

这也解释了@a_guest 指出的问题。需要第二个cumsum 调用以避免匹配第一个和第二个True 值之间的所有False 值。

---

如果性能很重要，请使用argmax 并设置值：

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 性能时序
我将借此机会炫耀perfplot，有一些时间，因为很高兴看到我们的解决方案如何随着不同大小的输入而变化。

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

趋势延续到更大的 N。cumsum 非常昂贵，而我的第二个解决方案和 @a_guest 的解决方案之间存在恒定的时间差。

Answer 2

您可以使用以下方法使用np.argmax 和带有np.logical_or.reduce 的产品来处理全部为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)

计时结果

提高性能的不同版本，即最快的方法排在最后：

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)