按列a分组后保留首行其余设为NaN，多列批量处理方案求助

原始数据 #

import pandas as pd

df = pd.DataFrame(
    {
        'a': [
            'a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b',
        ],
        'b': [
            -20, 20, 20, 20,-70, -70,-11, -100, -1, -1, -100, 100
        ],
        'c': [
            'f', 'f', 'f', 'f', 'f', 'x', 'x', 'k', 'k', 'k', 'k', 'k'
        ],
        'x': [
            'p', 'p', 'p', 'p', 'p', 'x', 'x', 'i', 'i', 'i', 'i', 'i'
        ],
    }
)

期望输出 #

a     b    c  x
0   a -20.0    f  p
1   a   NaN  NaN  p
2   a   NaN  NaN  p
3   a   NaN  NaN  p
4   a   NaN  NaN  p
5   a   NaN  NaN  x
6   b -11.0    x  x
7   b   NaN  NaN  i
8   b   NaN  NaN  i
9   b   NaN  NaN  i
10  b   NaN  NaN  i
11  b   NaN  NaN  i

需求说明 #

按列a分组，每个分组仅保留首行数据，其余行对应值改为NaN。实际场景中存在成百上千列，无法手动逐列指定。

原始尝试代码 #

import numpy as np 
import pandas as pd 

def func(g):
    g.iloc[1:, g.columns.get_loc('b')] = np.nan
    g.iloc[1:, g.columns.get_loc('c')] = np.nan
    return g

df = df.groupby('a', as_index=False).apply(func)

高效批量处理方案 #

不需要逐列指定，直接对除分组列a之外的所有列进行批量操作，代码如下：

import pandas as pd
import numpy as np

def keep_first_row(g):
    # 获取所有除分组列a之外的列
    target_cols = g.columns.difference(['a'])
    # 对每组第1行之后的目标列赋值NaN
    g.loc[g.index[1:], target_cols] = np.nan
    return g

# 分组处理，group_keys=False保持原始索引结构
result_df = df.groupby('a', group_keys=False).apply(keep_first_row)
print(result_df)

代码说明 #

• g.columns.difference(['a'])：自动筛选出所有非分组列，不管列数量多少都能适配，无需手动指定
• g.index[1:]：精准定位每组中除首行外的所有行
• group_keys=False：避免分组操作额外添加分组键索引，保证输出结构与原始DataFrame一致

运行上述代码后，即可得到符合期望的输出结果。

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