Boolean indexing array through array of boolean indexes without loop

Question

I want to index an array with a boolean mask through multiple boolean arrays without a loop.

This is what I want to achieve but without a loop and only with 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]

Divakar · Accepted Answer · 2019-06-13 14:18:43Z

6

Approach #1

Simply broadcast a to b's shape with np.broadcast_to and then mask it with b -

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

Approach #2

Another would be getting all the indices and mod those by the size of a, which would also be the size of each 2D block in b and then indexing into flattened a -

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

Approach #3

On the same lines as App#2, but keeping the 2D format and using non-zero indices along the last two axes of b -

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

Timings on large arrays (given sample shapes scaled up by 100x) -

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)

edited Jun 13, 2019 at 14:18

answered Jun 13, 2019 at 13:56

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Spark Monkay Over a year ago

In my code I had to do: np.expand_dims(a, axis=-1). The shape of my a is now (2, 3, 2) and the b is (2, 3).

Divakar Over a year ago

@SparkMonkay So, as inputs - data array a is 3D and boolean array/mask b is 2D?

Collectives™ on Stack Overflow

Boolean indexing array through array of boolean indexes without loop

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related