加速NumPy运算：优化移动窗口数组扩展的拼接性能

核心问题分析 #

频繁调用np.concatenate导致性能瓶颈的本质是：每次拼接都会触发内存重新分配与数据拷贝，百万级数据场景下多次重复操作会累积大量不必要的开销。解决思路是预先计算最终数组的完整形状，一次性分配内存后直接填充数据，彻底规避动态拼接的低效问题。

前提假设 #

假设原始数组df的字段顺序为：location（索引0）、var_id（索引1）、hour（索引2）、value（索引3），且我们需要生成的移动窗口为「包含当前小时在内的连续10个小时区间」。若数据未按hour排序，需先完成排序。

方案1：预分配内存+循环填充（低内存压力） #

import numpy as np

# 1. 按hour排序（若原始数据无序则执行）
df_sorted = df[df[:, 2].argsort()]

# 2. 提取关键字段与窗口参数
window_size = 10
unique_hours = np.unique(df_sorted[:, 2])
window_starts = unique_hours - (window_size - 1)

# 3. 用searchsorted快速定位每个窗口的起止索引（时间复杂度O(N log N)）
start_indices = np.searchsorted(df_sorted[:, 2], window_starts, side='left')
end_indices = np.searchsorted(df_sorted[:, 2], unique_hours, side='right')

# 4. 计算总数据量并预分配结果数组
window_lengths = end_indices - start_indices
total_rows = window_lengths.sum()
results = np.empty((total_rows, df_sorted.shape[1]), dtype=df_sorted.dtype)

# 5. 循环填充数据（仅内存赋值，无拷贝）
current_pos = 0
for start, end in zip(start_indices, end_indices):
    chunk_len = end - start
    results[current_pos:current_pos+chunk_len] = df_sorted[start:end]
    current_pos += chunk_len

方案2：向量化索引（极致性能） #

如果内存足够承载所有重复数据，可通过生成全局索引数组直接提取结果，完全避免循环：

import numpy as np

# 1. 排序与索引定位（同方案1）
df_sorted = df[df[:, 2].argsort()]
window_size = 10
unique_hours = np.unique(df_sorted[:, 2])
window_starts = unique_hours - (window_size - 1)
start_indices = np.searchsorted(df_sorted[:, 2], window_starts, side='left')
end_indices = np.searchsorted(df_sorted[:, 2], unique_hours, side='right')

# 2. 生成所有待提取的行索引
row_indices = np.concatenate([np.arange(s, e) for s, e in zip(start_indices, end_indices)])

# 3. 直接索引得到最终结果（仅一次内存拷贝）
results = df_sorted[row_indices]

额外优化建议 #

• 若需给每条数据标记所属的窗口目标小时，可预先生成重复的小时数组后拼接（仅一次np.hstack操作）
• 若数据量超大超出内存，可分批次处理窗口，每批次生成部分结果后写入磁盘，避免内存溢出
• 确保df的dtype统一且合理（比如用float32替代float64，int32替代int64），减少内存占用

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