如何在pandas DataFrame中按2x2滑动矩阵规则计算L2列

实现方案 #

直接通过numpy矩阵切片匹配2x2子矩阵的计算规则，逻辑清晰且运行高效：

import pandas as pd
import numpy as np

# 示例DataFrame
df = pd.DataFrame({'sales': ['2020-01','2020-02','2020-03','2020-04','2020-05','2020-06'],
                   '2020-01': [24,42,18,68,24,30],
                   '2020-02': [24,42,18,68,24,30],
                   '2020-03': [64,24,70,70,88,57],
                   '2020-04': [22,11,44,3,5,78],
                   '2020-05': [11,35,74,12,69,51]})

# 提取排除sales列的纯数值矩阵
val_arr = df.drop(columns='sales').values
n_rows, n_cols = val_arr.shape

# 初始化L2结果数组
l2 = np.zeros(n_rows, dtype=int)
# 首行按规则直接赋值
l2[0] = val_arr[0, 0]

# 遍历计算后续每个月的L2值
for i in range(1, n_rows):
    # 自动适配列边界，避免索引越界
    col_end = min(i + 1, n_cols)
    col_start = col_end - 2
    # 取左上角2x2子矩阵求和
    l2[i] = val_arr[i-1:i+1, col_start:col_end].sum()

# 结果写入DataFrame
df['L2'] = l2

print(df)

运行后得到的L2列值为[24, 132, 154, 187, 89, 203]，和预期输出完全匹配。

内容的提问来源于stack exchange，提问作者gyuyeong.cho