基于Pandas多层索引的上三角矩阵重塑：替代循环的优化方案问询

需求背景 #

需要对带有多层索引（第一层为DATE）的Pandas DataFrame执行上三角矩阵的重塑操作，当前采用循环逐个处理DATE分片的方式，希望找到更高效的实现方案。

示例数据与原实现 #

数据准备 #

import pandas as pd
import numpy as np

# 构造输入数据
A_1 = [1, .5, .3]
B_1 = [.5, 1, .4]
C_1 = [.3, .4, 1]

A_2 = [1, -.5, -.3]
B_2 = [-.5, 1, -.4]
C_2 = [-.3, -.4, 1]

# 构造带多层索引的DataFrame
DF_1 = pd.DataFrame({'A': A_1, 'B': B_1, 'C': C_1, 'DATE': '2024-05-23', 'ID': ['A', 'B', 'C']}).set_index(['DATE', 'ID'])
DF_2 = pd.DataFrame({'A': A_2, 'B': B_2, 'C': C_2, 'DATE': '2024-05-24', 'ID': ['A', 'B', 'C']}).set_index(['DATE', 'ID'])

DF = pd.concat([DF_1, DF_2])

原循环实现 #

def prior_solution(df: pd.DataFrame):
    # 单分片的上三角处理逻辑
    df_filt = df.where(np.triu(np.ones(df.shape)).astype(np.bool_))
    long = df_filt.stack().reset_index()
    long.columns = ['DATE', 'ID1', 'ID2', 'VALUE']
    return long

# 循环处理每个DATE分片
L_1 = prior_solution(DF_1)
L_2 = prior_solution(DF_2)

WANT = pd.concat([L_1, L_2])

更优实现方案 #

无需循环分片，可通过两种方式实现高效处理：

方案一：分组批量处理 #

利用groupby对DATE分组，批量执行上三角过滤与堆叠，保留逻辑简洁性：

def optimized_solution(df: pd.DataFrame):
    result = df.groupby(level='DATE').apply(
        lambda group: group.where(np.triu(np.ones(group.shape)).astype(bool)).stack().reset_index()
    ).reset_index(drop=True)
    result.columns = ['DATE', 'ID1', 'ID2', 'VALUE']
    return result

# 直接处理整个DF
OPTIMIZED_WANT = optimized_solution(DF)

方案二：向量化全局掩码 #

若所有DATE分片的结构一致，可构造全局掩码直接应用，避免分组开销：

def vectorized_solution(df: pd.DataFrame):
    # 获取单分片形状，生成基础上三角掩码
    single_shape = df.xs(df.index.get_level_values('DATE')[0]).shape
    base_mask = np.triu(np.ones(single_shape)).astype(bool)
    # 重复掩码匹配整个DataFrame的行数
    full_mask = np.tile(base_mask, (df.index.get_level_values('DATE').nunique(), 1))
    
    df_filt = df.where(full_mask)
    result = df_filt.stack().reset_index()
    result.columns = ['DATE', 'ID1', 'ID2', 'VALUE']
    return result

VECTORIZED_WANT = vectorized_solution(DF)

验证结果 #

两种优化方案的输出均与预期一致：

DATE ID1 ID2  VALUE
0  2024-05-23   A   A    1.0
1  2024-05-23   A   B    0.5
2  2024-05-23   A   C    0.3
3  2024-05-23   B   B    1.0
4  2024-05-23   B   C    0.4
5  2024-05-23   C   C    1.0
6  2024-05-24   A   A    1.0
7  2024-05-24   A   B   -0.5
8  2024-05-24   A   C   -0.3
9  2024-05-24   B   B    1.0
10 2024-05-24   B   C   -0.4
11 2024-05-24   C   C    1.0

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