numpy中的数组扩展

Question

我想做的是采用输入整数数组,并将其数据扩展为索引(例如,[2,1] - > [2,2,1]).如果术语关闭,我道歉 - 我不确定描述这个的最佳方式,因此,这可能是重复的.

以下是我现有方法的一个示例:

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

在嵌套的for循环中调用此方法,我希望从中挤出额外的性能.以下是当前的定时通话:

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

是否有可用于降低方法费用的不同方法？

Answer 1

本np.repeat应该做你最想要的东西:

a.repeat(a)

我的时间是5分钟你的88.

你的第一个例子是

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