Логический индексный массив через массив логических индексов без цикла

Я хочу проиндексировать массив с логической маской через несколько логических массивов без цикла.

Это то, чего я хочу добиться, но без цикла и только с numpy.

import numpy as np
a = np.array([[0, 1],[2, 3]])
b = np.array([[[1, 0], [1, 0]], [[0, 0], [1, 1]]], dtype=bool)

r = []
for x in b:
    print(a[x])
    r.extend(a[x])

# => array([0, 2])
# => array([2, 3])

print(r)
# => [0, 2, 2, 3]

# what I would like to do is something like this
r = some_fancy_indexing_magic_with_b_and_a
print(r)
# => [0, 2, 2, 3]

python numpy arrays matrix-indexing
person Spark Monkay    schedule 13.06.2019    source источник

Ответы (1)


arrow_upward
6
arrow_downward

Подход №1

Просто транслируйте форму a в b's с помощью np.broadcast_to, а затем замаскируйте его b -

In [15]: np.broadcast_to(a,b.shape)[b]
Out[15]: array([0, 2, 2, 3])

Подход № 2

Другим было бы получить все индексы и модифицировать их по размеру a, который также будет размером каждого блока 2D в b, а затем индексировать в сглаженный a -

a.ravel()[np.flatnonzero(b)%a.size]

Подход №3

В тех же строках, что и приложение № 2, но с сохранением формата 2D и использованием ненулевых индексов вдоль двух последних осей b -

_,r,c = np.nonzero(b)
out = a[r,c]

Тайминги на больших массивах (данные образцы форм увеличены в 100 раз) -

In [50]: np.random.seed(0)
    ...: a = np.random.rand(200,200)
    ...: b = np.random.rand(200,200,200)>0.5

In [51]: %timeit np.broadcast_to(a,b.shape)[b]
45.5 ms ± 381 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [52]: %timeit a.ravel()[np.flatnonzero(b)%a.size]
94.6 ms ± 1.64 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [53]: %%timeit
    ...: _,r,c = np.nonzero(b)
    ...: out = a[r,c]
128 ms ± 1.46 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
person Divakar    schedule 13.06.2019
comment
В моем коде мне нужно было сделать: np.expand_dims(a, axis=-1). Форма моего a теперь (2, 3, 2), а b - (2, 3). - person Spark Monkay; 13.06.2019
comment
@SparkMonkay Итак, в качестве входных данных - массив данных a равен 3D, а логический массив/маска b равен 2D? - person Divakar; 13.06.2019