如何基于多条件从Pandas DataFrame筛选行并处理特殊场景

原始DataFrame定义 #

import pandas as pd

df = pd.DataFrame({
    'user': ['user 1', 'user 4', 'user 1', 'user 4', 'user 1', 'user 4'],
    'group': [0, 0, 1, 1, 2, 2],
    'x1': [0.9, 0.9, 0.7, 0.7, 0.4, 0.4],
    'x2': [0.759740, 1.106061, 0.619357, 1.260234, 0.540633, 1.437956]
})

原始输出 #

user  group   x1        x2
0  user 1      0  0.9  0.759740
1  user 4      0  0.9  1.106061
2  user 1      1  0.7  0.619357
3  user 4      1  0.7  1.260234
4  user 1      2  0.4  0.540633
5  user 4      2  0.4  1.437956

需求说明 #

• 针对每个user：
  1. 若存在x2 < x1的行，选择其中x2值最小的那一行
  2. 若不存在x2 < x1的行，选择该user所有行中x2最小的行，并将该行的group字段修改为10

示例说明 #

• user 1有多个满足x2 < x1的行，其中第2行的x2值最小，因此选择该行
• user 4所有行的x2均大于x1，因此选择其x2最小的第1行，并将group改为10

解决方案代码 #

def process_user(group):
    # 筛选当前user下x2小于x1的行
    valid_subset = group[group['x2'] < group['x1']]
    if not valid_subset.empty:
        # 返回x2最小的行
        return valid_subset.loc[valid_subset['x2'].idxmin()]
    else:
        # 取x2最小的行，修改group为10后返回
        min_row = group.loc[group['x2'].idxmin()].copy()
        min_row['group'] = 10
        return min_row

# 按user分组处理并生成结果
result_df = df.groupby('user', group_keys=False).apply(process_user).reset_index(drop=True)

预期输出 #

user  group   x1        x2
0  user 1      1  0.7  0.619357
1  user 4     10  0.9  1.106061

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