Search code examples
pythonarrayslistnumpyvectorization

Numpy create array from list of integers specifying the array elements

Say we have the list a = [3,4,2,1], from a we want to obtain the following array: b = [0,0,0,1,1,1,1,2,2,3].

I have managed to do it like this:

import numpy as np 
a = [3,4,2,1]
b = np.concatenate([[i] * n for i, n in enumerate(a])
print(b)

Output:

array([0, 0, 0, 1, 1, 1, 1, 2, 2, 3])

Which works fine, but I can't help but wonder if there is a better way?

EDIT @mathfux's solution is much more elegant than mine by simply using np.repeat. I benchmarked the two using perfplot, here the results: enter image description here Benchmark code:

import numpy as np
import perfplot

out = perfplot.bench(
    setup=lambda n: np.arange(n),
    kernels=[
        lambda x: np.concatenate([[i] * n for i, n in enumerate(x)]),
        lambda x: np.repeat(range(len(x)), x)
    ],
    n_range=np.arange(128, 128 * 16, 128),
    labels=[
        "concat",
        "repeat"
    ],
    xlabel='np.arange(x)'
)
out.show()
Solution

  • Try np.repeat: np.repeat([0,1,2,3], [3,4,2,1])