CODEWITHSUNDEEP
pythonnumpy
John Makkonen
John Makkonen

Reputation: 43

Element-wise boolean multiplication

I'm looking for the most beautiful and short way to multiply those types of lists:

a = [True, False, True, False, False]

b = [100, 200]

Length of b equals to number of True elements in a

The answer that I need is [100, 0, 200, 0, 0] here

Are there any simple ways to get that answer?
It looks like element-wise multiplication, but the fact is that the size of the second list is smaller, so common numpy methods don't work without bad code.

Hope, you'll find nice solution

Upvotes: 4

Views: 4126

Answers (4)

Daniel F
Daniel F

Reputation: 14399

Swiping a bit from @AGN above

c = np.zeros_like(a, dtype = b.dtype)
c[a] = b

Upvotes: 1

sacuL
sacuL

Reputation: 51395

You could do this in numpy:

c = np.array(a).astype(int)

c[c==1] = b

>>> c
array([100,   0, 200,   0,   0])

Note If you need the result as a list (based on your desired output), rather than a numpy array, use c.tolist()

Upvotes: 10

AGN Gazer
AGN Gazer

Reputation: 8378

Another option is to use either numpy.place() or numpy.put():

c = np.zeros_like(a, dtype=np.result_type(b))
np.place(c, a, b)

or

c = np.zeros_like(a, dtype=np.result_type(b))
np.put(c, np.flatnonzero(a), b)

I prefer the first option.

Timings:

Generate larger data set:

In [39]: a = np.random.randint(0, 2, 100000).astype(np.bool_)

In [40]: b = np.random.randint(0, 100000, np.sum(a))

@sacul solution:

In [41]: %timeit c = np.array(a).astype(int); c[c==1] = b
621 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

My first option:

In [42]: %timeit c = np.zeros_like(a, dtype=np.result_type(b)); np.place(c, a, b)
462 µs ± 8.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Upvotes: 2

Patrick Haugh
Patrick Haugh

Reputation: 61032

There's a pretty clean iterator based solution

it = iter(b)
[next(it) if x else 0 for x in a]
# [100, 0, 200, 0, 0]

Upvotes: 7

Related Questions