CODEWITHSUNDEEP
pythonnumpy
hetsch
hetsch

Reputation: 1568

Expand a numpy array based on a value in another array

I have the following numpy array a = np.array([1,1,2,1,3]) that should be transformed into the following array b = np.array([1,1,1,1,1,1,1,1]).

What happens is that all the non 1 values in the a array should be expanded in the b array to their multiple defined in the a array. Simpler said, the 2 should become 2 ones, and the 3 should become 3 ones.

Frankly, I couldn't find a numpy function that does this, but I'm sure one exists. Any advice would be very welcome! Thank you!

Upvotes: 2

Views: 266

Answers (3)

hpaulj
hpaulj

Reputation: 231355

In [71]: np.ones(len(a),int).repeat(a)                                          
Out[71]: array([1, 1, 1, 1, 1, 1, 1, 1])

For this small example it is faster than np.ones(a.sum(),int), but it doesn't scale quite as well. But overall both are fast.

Upvotes: 1

Divakar
Divakar

Reputation: 221524

We can simply do -

np.ones(a.sum(),dtype=int)

This will accomodate all numbers : 1s and non-1s, because of the summing and hence give us the desired output.

Upvotes: 2

Kasravnd
Kasravnd

Reputation: 107287

Here's one possible way based on the number you wanna be repeated:

In [12]: a = np.array([1,1,2,1,3])
In [13]: mask = a != 1
In [14]: np.concatenate((a[~mask], np.repeat(1, np.prod(a[mask]))))
Out[14]: array([1, 1, 1, 1, 1, 1, 1, 1, 1])

Upvotes: 0

Related Questions