如何按自定义偏好规则移除Pandas中的重复数据？

问题描述 #

Pandas的drop_duplicates方法仅支持保留重复项的第一个、最后一个或全部删除，但需要实现更复杂的去重逻辑：

• 当检测到基于id和val列的重复项时，若其中一条的tag列值属于指定偏好集合，保留该条（不受位置限制）
• 若多条重复项的tag都属于偏好集合，按默认drop_duplicates逻辑（如保留第一个）处理
• 若所有重复项的tag都不属于偏好集合，同样按默认逻辑处理

尝试用掩码实现但结果不符合预期，代码及问题如下：

import pandas as pd

def conditional_remove_duplicates(df, preferred_tags):
    duplicates_mask = df.duplicated(subset=['id', 'val'], keep=False)
    preferred_mask = df['tag'].isin(preferred_tags)
    mask = duplicates_mask | preferred_mask
    df = df[mask].drop_duplicates(subset=['id', 'val'], keep='first')
    return df

data = {'id': ['A', 'A', 'A', 'A', 'B', 'B', 'C', 'D', 'D'],
        'val': [10, 10, 11, 10, 20, 20, 30, 40, 40],
        'tag': ['X', 'Z', 'X', 'Y', 'Z', 'X', 'X', 'Z', 'Z']}
preferred_tags = {'X', 'Y'}

df = pd.DataFrame(data)
print(df)
"""
  id  val tag
0  A   10   X
1  A   10   Z
2  A   11   X
3  A   10   Y
4  B   20   Z
5  B   20   X
6  C   30   X
7  D   40   Z
8  D   40   Z
"""
result_df = conditional_remove_duplicates(df, preferred_tags)
print(result_df)
""" 实际输出：
  id  val tag
0  A   10   X
2  A   11   X
4  B   20   Z
6  C   30   X
7  D   40   Z

期望输出：
   id  val tag
0  A   10   X
2  A   11   X
5  B   20   X
6  C   30   X
7  D   40   Z
"""

解决方案 #

核心思路是通过标记优先级、排序后再去重，确保偏好项被优先保留：

1. 给tag属于偏好集合的行设置高优先级，其他行设置低优先级
2. 按分组键、优先级、原索引排序，保证偏好项排在同组最前，同优先级时保留最早的行
3. 最后对分组后的结果去重，保留每组的第一行

修改后的代码如下：

import pandas as pd

def conditional_remove_duplicates(df, preferred_tags):
    # 添加优先级列：偏好tag标记为1，其他为0
    df['priority'] = df['tag'].isin(preferred_tags).astype(int)
    # 按分组键、优先级降序、原索引升序排序
    sorted_df = df.sort_values(
        by=['id', 'val', 'priority', df.index],
        ascending=[True, True, False, True]
    )
    # 去重并移除临时优先级列
    result_df = sorted_df.drop_duplicates(subset=['id', 'val'], keep='first').drop(columns='priority')
    # 恢复原索引顺序（可选，按需保留）
    result_df = result_df.sort_index()
    return result_df

data = {'id': ['A', 'A', 'A', 'A', 'B', 'B', 'C', 'D', 'D'],
        'val': [10, 10, 11, 10, 20, 20, 30, 40, 40],
        'tag': ['X', 'Z', 'X', 'Y', 'Z', 'X', 'X', 'Z', 'Z']}
preferred_tags = {'X', 'Y'}

df = pd.DataFrame(data)
result_df = conditional_remove_duplicates(df, preferred_tags)
print(result_df)

代码说明 #

• 优先级标记：将布尔值转换为整数，让偏好项的优先级高于非偏好项
• 排序逻辑：先按id和val分组，再按优先级降序确保偏好项在前，最后按原索引升序保证同优先级时保留最早的行
• 去重处理：此时每组的第一行就是符合要求的目标行，直接去重即可
• 索引恢复：可选步骤，若不需要保持原数据顺序可省略

运行后输出与期望结果一致：

id  val tag
0  A   10   X
2  A   11   X
5  B   20   X
6  C   30   X
7  D   40   Z

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