Как я могу получить индексы массива при некоторых условиях

Question

Я сталкиваюсь с такой проблемой: предположим, у меня есть массивы, подобные этой: a = np.array([[1,2,3,4,5,4,3,2,1],]) label = np.array([[1,0,1,0,0,1,1,0,1],]) Мне нужно получить индексы a, в этой позиции значение элемента label равно 1, а значение aнаибольшее значение, из-за которого label равно 1.

. В приведенном выше примере может показаться странным, что индексы, где label равен 1: 0, 2, 5, 6, 8, ихсоответствующие значения a, таким образом: 1, 3, 4, 3, 1, среди которых 4 - значения, поэтому мне нужно получить результат 5, который является индексом числа 4 в a.Как я мог сделать это с NumPy?

Answer 1

Вы можете использовать маскированные массивы:

>>> np.ma.masked_where(~label.astype(bool), a).argmax()
5

Answer 2

Вот один из самых простых способов.

>>> np.argmax(a * (label == 1))
5
>>> np.argmax(a * (label == 1), axis=1)
array([5])

Метод Coldspeed может занять больше времени.

Answer 3

Получите индексы 1s, скажем, как idx, затем внесите в индекс a, получите индекс max и, наконец, отследите его до первоначального порядка, указав в idx -

idx = np.flatnonzero(label==1)
out = idx[a[idx].argmax()]

Пример выполнения -

# Assuming inputs to be 1D
In [18]: a
Out[18]: array([1, 2, 3, 4, 5, 4, 3, 2, 1])

In [19]: label
Out[19]: array([1, 0, 1, 0, 0, 1, 1, 0, 1])

In [20]: idx = np.flatnonzero(label==1)

In [21]: idx[a[idx].argmax()]
Out[21]: 5

Для a в виде целых чисел и label в виде массива 0s и 1s, мы могли бы оптимизировать дальше, поскольку мы могли бы масштабировать на основе aв диапазоне значений в нем, например, так:

(label*(a.max()-a.min()+1) + a).argmax()

Кроме того, если a имеет только положительные числа, это упрощается до -

(label*(a.max()+1) + a).argmax()

Время для положительных целых чиселa -

In [115]: np.random.seed(0)
     ...: a = np.random.randint(0,10,(100000))
     ...: label = np.random.randint(0,2,(100000))

In [117]: %%timeit
     ...: idx = np.flatnonzero(label==1)
     ...: out = idx[a[idx].argmax()]
1000 loops, best of 3: 592 µs per loop

In [116]: %timeit (label*(a.max()-a.min()+1) + a).argmax()
1000 loops, best of 3: 357 µs per loop

# @coldspeed's soln
In [120]: %timeit np.ma.masked_where(~label.astype(bool), a).argmax()
1000 loops, best of 3: 1.63 ms per loop

# won't work with negative numbers in a
In [119]: %timeit (label*(a.max()+1) + a).argmax()
1000 loops, best of 3: 292 µs per loop

# @klim's soln (won't work with negative numbers in a)
In [121]: %timeit np.argmax(a * (label == 1))
1000 loops, best of 3: 229 µs per loop