python - Array expansion in numpy -

what i'd take input integer array, , expand data indices (e.g., [2, 1] -> [2, 2, 1]). apologize if terminology off -- wasn't sure of best way describe this, such, possible duplicate.

here example of current method have in place:

>>> def expand(a): ...     b = np.empty(a.sum(), dtype=np.int32) ...     idx = 0 ...     in a: ...         j in range(i): ...             b[idx] = ...             idx += 1 ...     return b ...  >>> = np.array([3, 2, 1, 4]) >>> expand(a) array([3, 3, 3, 2, 2, 1, 4, 4, 4, 4], dtype=int32) 

this method called within nested loop i'd squeeze additional performance out of. below current timing call:

>>> = np.random.randint(0, 1000, 1000) >>> %timeit expand(a) 10 loops, best of 3: 86.9 ms per loop  

is there different approach used lower expense of method?

the np.repeat should of want:

a.repeat(a) 

i timeit @ 5ms v 88.

your first example be

arange(2).repeat([2,1])