1

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()

1 Answer 1

3

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

Sign up to request clarification or add additional context in comments.

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.