本文介绍如何在 pandas 中高效计算“半对角线”（即从左上到右下、逐行偏移的斜向序列）的平均值，适用于带下三角 nan 的稀疏矩阵结构，核心是利用 numpy 的 np.take_along_axis 与 np.nanmean 实现无循环向量化操作。
本文介绍如何在 pandas 中高效计算“半对角线”（即从左上到右下、逐行偏移的斜向序列）的平均值，适用于带下三角 nan 的稀疏矩阵结构，核心是利用 numpy 的 np.take_along_axis 与 np.nanmean 实现无循环向量化操作。
在实际数据分析中，我们常遇到类似“对角线填充”的表格结构：有效数据沿主对角线及其右上方分布，左下方为缺失值（NaN），例如实验时间序列、递推计算结果或状态转移矩阵。此时，若需为每一行提取其对应的“斜向序列”（如第 0 行取列 1~18，第 1 行取列 2~18，依此类推），并求该序列的平均值，传统循环或 apply 效率低下且不易扩展。
Pandas 本身不提供原生对角线分组聚合函数，但可借助 NumPy 的底层数组操作实现完全向量化的解决方案。关键思路是：将目标数值区域（本例中为第 3 列起的所有数字列）转为 NumPy 数组，然后对每行按 NaN 位置排序，把非空值“左对齐”，最后按列求均值——这恰好对应每条斜向序列的平均值。
以下是完整实现：
import pandas as pd
import numpy as np
# 构建示例 DataFrame（略去 data 字典定义，见原文）
df = pd.DataFrame(data)
# ✅ 方法一：适用于标准下三角 NaN 结构（推荐）
numeric_cols = df.columns[3:] # 提取所有数值列（1, 2, ..., 18）
a = df[numeric_cols].to_numpy() # 转为 (n_rows, n_cols) 数组
# 对每行，按是否为 NaN 排序（NaN 排在末尾），实现“非空值左对齐”
aligned = np.take_along_axis(
a,
np.argsort(np.isnan(a), axis=1, kind='stable'),
axis=1
)
# 沿行方向（axis=0）对齐后矩阵求列均值（自动跳过 NaN）
diag_avg = np.nanmean(aligned, axis=0)
# 将结果作为新列添加回 DataFrame
df['diag_avg'] = diag_avg
? 为什么有效？
原始数据中，第 i 行的有效值从第 i+1 列开始（0-indexed）。np.argsort(np.isnan(a), axis=1) 会为每行生成一个索引排列，使 False（非 NaN）排在 True（NaN）之前，从而将每行的“有效斜向片段”统一左对齐至前若干列。随后 np.nanmean(..., axis=0) 即对齐后第 j 列的所有值求均值——而这列恰好汇集了所有起始于列 j 的斜向元素（即原问题中 SP=0 对应列 1 的所有值、SP=1 对应列 2 的所有值……），完美匹配业务语义。
⚠️ 注意事项：
- 若数据中存在上三角 NaN（即有效值不连续、中间有空缺），应改用更鲁棒的掩码排序法：
m = ~df[numeric_cols].notna().cummax(axis=1).to_numpy() # True for positions before first NaN in row aligned = np.take_along_axis(a, np.argsort(m, axis=1, kind='stable'), axis=1) df['diag_avg'] = np.nanmean(aligned, axis=0)
- 确保目标列为数值类型；若含字符串或混合类型，先用 df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce') 清洗。
- 该方法时间复杂度为 O(nm log m)（m 为列数），远优于 Python 循环的 O(nm) 实际开销，尤其适合万级行以上场景。
最终，df['diag_avg'] 即为所求的“半对角线平均值”列，其第 i 个值代表 SP=i 对应斜向序列的均值，结果精确、高效且可复用。
