Как получить только первое значение True из каждой строки массива numpy?

Question

У меня есть логический массив nxy размером 4x3, и я пытаюсь вернуть массив одинакового размера со значением False, за исключением расположения первого значения True в каждой строке оригинала. Так что, если у меня есть начальный массив

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

Вы можете использовать следующий подход, используя 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)

Answer 2

Вы можете использовать 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 необходим, чтобы избежать совпадения всех значений False между первым и вторым значением True.

---

Если важна производительность, используйте 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's есть постоянная разница во времени.