Skip to main content
 首页 » 编程设计

python之numpy 中的数组扩展

2025年05月04日10jackei

我想做的是获取一个输入整数数组,并将其数据扩展为索引(例如,[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  

是否可以使用不同的方法来降低该方法的费用?

请您参考如下方法:

np.repeat 应该可以完成您想要的大部分操作:

a.repeat(a) 

我计时 5 毫秒,而你的 88。

你的第一个例子是

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