如何高效创建特定重复模式的NumPy整数数组？

问题回顾 #

需要生成一个NumPy整数数组，其中整数n恰好重复n次（从1到指定最大值max），例如max=4时结果为[1,2,2,3,3,3,4,4,4,4]。原有的循环+np.append方法在max较大时性能极差，原因是每次np.append都会重新分配内存并复制现有数据，时间复杂度为O(n²)。

高效实现方案 #

方案1：使用np.repeat（最简洁高效）

利用NumPy的向量化操作np.repeat，直接生成目标数组，底层为C级优化，时间复杂度O(n)：

import numpy as np

max_val = 4
my_arr = np.repeat(np.arange(1, max_val + 1), np.arange(1, max_val + 1))
print(my_arr)
# 输出：array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])

方案2：预分配内存后切片赋值

如果需要手动控制内存分配，可以先计算总长度，一次性分配数组空间，再通过循环切片赋值，避免重复内存复制：

import numpy as np

max_val = 4
total_length = max_val * (max_val + 1) // 2  # 等差数列求和公式计算总长度
my_arr = np.empty(total_length, dtype=np.int64)

start_idx = 0
for i in range(1, max_val + 1):
    end_idx = start_idx + i
    my_arr[start_idx:end_idx] = i
    start_idx = end_idx

print(my_arr)
# 输出：array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])

性能对比 #

当max=1000时，np.repeat方法的执行速度是原循环append方法的数百倍，内存占用也更低（无需多次创建临时数组）。

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