YXD
YXD

Reputation: 32521

Retrieving array elements with an array of frequencies in NumPy

I have an array of numbers, a. I have a second array, b, specifying how many times I want to retrieve the corresponding element in a. How can this be achieved? The ordering of the output is not important in this case.

import numpy as np

a = np.arange(5)
b = np.array([1,0,3,2,0])

# desired output = [0,2,2,2,3,3]
# i.e. [a[0], a[2], a[2], a[2], a[3], a[3] ]

Upvotes: 2

Views: 152

Answers (3)

none
none

Reputation: 12087

here's one way to do it:

res = []
for i in xrange(len(b)):
    for j in xrange(b[i]):
        out.append(a[i])

res = np.array(res)  # optional

Upvotes: 0

seberg
seberg

Reputation: 8975

Thats exactly what np.arange(5).repeat([1,0,3,2,0]) does.

Upvotes: 6

jlengrand
jlengrand

Reputation: 12817

A really inefficient way to do that is this one :

import numpy as np

a = np.arange(5)
b = np.array([1,0,3,2,0])

res = []
i = 0
for val in b:
    for aa in range(val):
        res.append(a[i])
    i += 1
print res

Upvotes: 0

Related Questions