如何利用Numpy内置函数高效生成指定区间重复元素的数组？

给定数组a = np.asarray([0,1,2])和b = np.asarray([3,7,10])，要生成数组c（其中元素0重复3次、1重复4次、2重复3次，对应b中相邻索引的差值），完全可以用Numpy内置的np.repeat()函数高效实现，比Python循环快得多。

实现步骤 #

1. 计算每个元素的重复次数：
   观察b的索引位置可知，第一个元素0的重复次数是b[0] - 0 = 3，第二个元素1的次数是b[1]-b[0] = 4，第三个元素2的次数是b[2]-b[1] = 3。可以用np.diff()结合拼接初始值的方式批量计算：

   counts = np.diff(np.concatenate([[0], b]))
   

   得到的counts结果为array([3,4,3])。

2. 生成目标数组：

   import numpy as np
   
   a = np.asarray([0,1,2])
   b = np.asarray([3,7,10])
   
   counts = np.diff(np.concatenate([[0], b]))
   c = np.repeat(a, counts)
   

验证结果 #

执行后c的值为：

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

完全符合需求。

效率说明 #

np.repeat()是Numpy底层优化的向量化操作，直接在C层面处理数组，避免了Python循环的额外开销，数据量越大，效率优势越明显。

内容的提问来源于stack exchange，提问作者McM