如何压缩含缺失值的多层索引pd.DataFrame行数据？

问题场景 #

给定一个包含多层索引和缺失值的Pandas DataFrame，需要将具有相同索引的行合并，填充缺失值以压缩数据。

原始代码 #

import numpy as np
import pandas as pd

# 创建多层索引
index = pd.MultiIndex.from_tuples(
    [
        ("A", "X", "I"),
        ("A", "X", "I"),
        ("A", "Y", "I"),
        ("A", "Y", "II"),
        ("A", "Y", "I"),
    ],
    names=["level_1", "level_2", "level_3"],
)

# 创建DataFrame
data = [[1, np.nan], [np.nan, 1], [np.nan, 1], [np.nan, 1], [1, np.nan]]
df = pd.DataFrame(data, index=index, columns=["column1", "column2"])

print(df)

原始输出 #

column1  column2
level_1 level_2 level_3                  
A       X       I            1.0      NaN
                I            NaN      1.0
        Y       I            NaN      1.0
                II           NaN      1.0
                I            1.0      NaN

期望压缩结果 #

column1  column2
level_1 level_2 level_3                  
A       X       I            1.0      1.0
        Y       I            1.0      1.0
                II           NaN      1.0

解决方法 #

通过groupby按多层索引分组，再使用max()方法聚合即可实现需求——max()会自动忽略NaN值，保留每组内各列的有效数值，完美合并重复索引行并填充缺失值。

实现代码 #

# 按多层索引分组并聚合
compressed_df = df.groupby(level=["level_1", "level_2", "level_3"]).max()

print(compressed_df)

补充说明 #

如果同一索引组的同一列存在多个非NaN值，max()会取最大值；若需保留第一个出现的非NaN值，可改用first()方法（需确保每组内同列非NaN值唯一，否则会丢失后续值）：

compressed_df = df.groupby(level=["level_1", "level_2", "level_3"]).first()

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