Efficiently invert floats in a list based on a con

Question

I have one ndarray of floats (a) and one list of 0 or 1 (b). a and b have the same length, typically 1000. I would like to apply a simple function (such as minus the inverse) to the elements of a whose index correspond to 1 in b.

The following method takes a bit less than 1 millisecond. Would it be possible to make it faster?

import numpy as np

# generate dummy data
n = 1000
a = np.random.uniform(low=0, high=10, size=n)
b = np.random.choice(2, 1000)

# map the function
start = time.time()
out = [a[i] if b[i] == 0 else -1/a[i] for i in range(n)]
print time.time() - start

Answer 1

Use b as a mask, and set a's cells accordingly:

m = np.array(b).astype(bool)
a[m] = -1 / a[m]

Even better, initialise b using np.random.choice:

b = np.random.choice(2, 1000).astype(bool)

Now, you don't need the overhead of converting b to an array, just use it directly to index a:

a[b] = -1 / a[b]

This runs in

22.3 µs ± 501 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)