Python Pandas：长格式数据按Selector变量保留观测值及转宽格式

原始数据（存储于Pandas DataFrame）：

ID  Position    Select
1      A          0
2      B          1
2      C          0
3      B          0
3      C          0
4      A          1
5      A          0

需求说明 #

1. 筛选规则：
   • 保留所有单行且Select=1的ID行；
   • 若某ID对应多行且任意一行Select=1，则保留该ID的所有行；
     筛选后中间结果：

   ID  Position    Select
   2      B          1
   2      C          0
   4      A          1
   

2. 将筛选后的数据转换为宽格式，最终结果：

   ID  Position1   Position2  Select
   2      B            C        1
   4      A                     1
   

实现方案 #

第一步：数据筛选

最优方式是利用groupby.transform标记符合条件的ID，再进行布尔索引筛选，代码简洁且效率高：

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'ID': [1,2,2,3,3,4,5],
    'Position': ['A','B','C','B','C','A','A'],
    'Select': [0,1,0,0,0,1,0]
})

# 为每一行标记对应ID是否存在Select=1的记录
has_valid_select = df.groupby('ID')['Select'].transform(lambda x: x.any())
# 直接筛选符合条件的行
filtered_df = df[has_valid_select]

解释：transform会返回与原数据同长度的结果，自动覆盖了单行/多行的两种场景——单行且Select=1时，x.any()返回True；多行且存在Select=1的ID，所有行都会被标记为True。

第二步：转换为宽格式

通过为每个ID的Position生成带序号的列名，再用pivot实现长转宽：

# 为每个ID下的Position生成序号后缀（Position1、Position2...）
filtered_df['pos_seq'] = filtered_df.groupby('ID').cumcount() + 1
filtered_df['pos_col'] = 'Position' + filtered_df['pos_seq'].astype(str)

# 执行长转宽操作
wide_df = filtered_df.pivot(
    index='ID',
    columns='pos_col',
    values='Position'
).reset_index()

# 添加Select列（取每个ID的Select最大值，确保只要有1就保留1）
wide_df['Select'] = filtered_df.groupby('ID')['Select'].max()

# 调整列顺序并填充空值为空白
wide_df = wide_df[['ID', 'Position1', 'Position2', 'Select']].fillna('')

最终wide_df即为目标宽格式数据。

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