如何在Pandas中无需脆弱变通方法保留分组列

针对你提到的两个痛点，以下是基于Pandas原生API的优雅解决方法：

问题1：避免用g.name[0]这类按位置访问分组键的脆弱写法 #

当多列分组时，g.name返回的是对应分组值的元组，按索引取值（如g.name[0]）依赖分组列的顺序，一旦分组列顺序调整就会出错。推荐两种更健壮的方式：

方式1：映射分组键与列名 #

将分组键元组与列名绑定为字典，直接通过列名获取对应值：

.groupby(["group1", "group2"])
.apply(lambda g: 
    g.assign(
        flag=lambda df: (
            False if dict(zip(g.index.names, g.name))['group1'] == "a" 
            else ~(rng.random(len(g)) < 0.5)
        )
    )
)

方式2：直接取组内分组列的固定值 #

同一分组内的分组列值完全一致，可直接取组内该列的第一个值：

.groupby(["group1", "group2"])
.apply(lambda g: 
    g.assign(
        flag=lambda df: (
            False if g['group1'].iloc[0] == "a" 
            else ~(rng.random(len(g)) < 0.5)
        )
    )
)

问题2：避免apply后分组列转为索引，无需reset_index() #

Pandas 2.1.0及以上版本新增了groupby.apply的include_groups参数，设置为True时，分组列会直接保留在结果的列中，而非成为索引的一部分，无需手动重置索引。

优化后的完整代码 #

import pandas as pd
import numpy as np

print("pandas version:", pd.__version__)
# pandas version: 3.0.1
print("numpy version:", np.__version__)
# numpy version: 2.4.2

rng = np.random.default_rng(5)

df = pd.DataFrame({
    "group1": ["a", "a", "b", "b", "b"],
    "group2": ["X", "X", "Y", "Y", "Y"],
    "value": [1, 2, 3, 4, 5]
})

df2 = (
    df
    .groupby(["group1", "group2"], include_groups=True)
    .apply(lambda g: g.assign(
        flag=lambda df: (
            False if g['group1'].iloc[0] == "a" else ~(rng.random(len(g)) < 0.5)
        )
    ))
)

print(df2)

输出结果 #

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