利用Numpy/Pandas广播机制实现大二维数组列向滑动窗口均值

问题背景 #

现有10000×10000的大型二维Numpy数组（或Pandas DataFrame），需要将数组内元素替换为对应列中最近n行的均值。以n=3为例，具体规则如下：

原始数组 #

[[10, 30, 8, 1],
 [11, 5, 19, 12],
 [12, 18, 15, 6],
 [13, 10, 21, 9],
 [14, 67, 14, 2],
 [15, 13, 12, 6]]

处理后结果（n=3） #

[[12.5, 23.5, 14.83333333, 5.833333333],
 [12, 10.33333333, 18.33333333, 9],
 [13, 31.66666667, 16.66666667, 5.666666667],
 [14, 30, 15.66666667, 5.333333333]]

规则说明 #

• 结果中的每个元素，对应原始数组同一列中最近n行的均值（比如结果里的14是原始列中15、14、13的平均值）
• 当前用双层Python循环实现效率极低，需要基于Numpy/Pandas的向量化操作优化。

解决方案 #

1. Numpy 实现（高性能推荐） #

利用Numpy的步长技巧生成滑动窗口，全程向量化计算，无Python循环，内存占用低，速度极快。

import numpy as np

def rolling_col_mean_numpy(arr, n):
    rows, cols = arr.shape
    # 获取数组的内存步长
    row_stride, col_stride = arr.strides
    # 生成滑动窗口视图（不复制数据）
    window_view = np.lib.stride_tricks.as_strided(
        arr,
        shape=(rows - n + 1, n, cols),
        strides=(row_stride, row_stride, col_stride),
        writeable=False  # 避免意外修改原始数组
    )
    # 沿窗口维度计算均值
    return window_view.mean(axis=1)

# 测试示例
original_arr = np.array([
    [10, 30, 8, 1],
    [11, 5, 19, 12],
    [12, 18, 15, 6],
    [13, 10, 21, 9],
    [14, 67, 14, 2],
    [15, 13, 12, 6]
])

print(rolling_col_mean_numpy(original_arr, n=3))

关键点：as_strided通过修改数组的步长信息创建滑动窗口视图，不复制原始数据，这是处理大数组时的核心优势。

2. Pandas 实现（简洁易用） #

Pandas的rolling方法封装了滚动窗口逻辑，语法更直观，适合已经用Pandas处理数据的场景。

import pandas as pd

def rolling_col_mean_pandas(df, n):
    # 按列计算滚动均值，窗口大小n，仅保留有完整n个元素的窗口结果
    rolling_result = df.rolling(window=n, axis=0, min_periods=n).mean()
    # 移除前n-1个无效的NaN行，得到最终结果
    return rolling_result.dropna().values

# 测试示例
original_df = pd.DataFrame([
    [10, 30, 8, 1],
    [11, 5, 19, 12],
    [12, 18, 15, 6],
    [13, 10, 21, 9],
    [14, 67, 14, 2],
    [15, 13, 12, 6]
])

print(rolling_col_mean_pandas(original_df, n=3))

关键点：rolling的axis=0指定按行滚动（即针对列计算），min_periods=n确保只有窗口填满n个元素时才计算均值，最后dropna()过滤掉无效行。

内容的提问来源于stack exchange，提问作者Jan Vaško