布尔索引数组到不带循环的布尔索引数组

Question

我想通过多个布尔数组而不是循环来索引带有布尔掩码的数组。

这是我想要实现的，但没有循环，只有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]

Answer 1

方法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)