基于多条件筛选DataFrame行，不满足条件则分配随机新组号

原始数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'user': ['user 1', 'user 2', 'user 3', 'user 4', 'user 1', 'user 2', 'user 3', 'user 4'],
    'group': [0, 0, 0, 0, 1, 1, 1, 1],
    'p1': [0.759, 1.106, 1.619, 1.260, 0.540, 1.437, 1.440, 1.332],
    'p2': [0.9, 0.9, 0.9, 0.9, 0.7, 0.7, 0.7, 0.7],
})

需求说明 #

• 对每个用户：
  1. 筛选出满足p1 < p2的行，若存在此类行，选取其中p1值最小的行对应的group号
  2. 若所有行均不满足p1 < p2，分配一个不在现有group列表中的随机组号

解决方案代码 #

# 标记满足条件的行
df['valid'] = df['p1'] < df['p2']

# 提取有有效行的用户，取p1最小的对应group
valid_users_result = df[df['valid']].sort_values(['user', 'p1'])\
                          .groupby('user')['group'].first().reset_index()

# 找出无有效行的用户
all_users = df['user'].unique()
invalid_users = list(set(all_users) - set(valid_users_result['user']))

# 生成不在现有group中的随机组号（确保不重复）
existing_groups = df['group'].unique()
max_existing_group = existing_groups.max()
# 生成足够范围的随机数，避免重复
new_group_numbers = np.random.choice(
    range(max_existing_group + 1, max_existing_group + len(invalid_users)*2),
    size=len(invalid_users),
    replace=False
)

# 构造无有效行用户的结果
invalid_users_result = pd.DataFrame({
    'user': invalid_users,
    'group': new_group_numbers
})

# 合并最终结果
final_result = pd.concat([valid_users_result, invalid_users_result], ignore_index=True)

print(final_result)

代码说明 #

1. 标记有效行：新增valid列快速识别满足p1 < p2的行
2. 处理有效用户：筛选有效行后按用户和p1升序排序，每个用户取第一行的group（即p1最小的有效行对应组）
3. 识别无效用户：找出没有任何有效行的用户集合
4. 生成随机新组号：基于现有最大组号生成范围足够的随机数，用replace=False保证每个新组号唯一且不在现有组列表中
5. 合并结果：将有效用户和无效用户的结果合并，得到最终每个用户的组号分配

示例输出（随机组号部分会变化） #

user  group
0  user 1      1
1  user 2      2
2  user 3      4
3  user 4      3

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