用更复杂的优先级去除Pandas重复项

Question

Pandas的.drop_duplicates功能可以让我们选择保留重复数据中的第一个、最后一个，或者不保留任何重复项。但我遇到了一个更复杂的情况。假设我有一组首选值，如果发现一对重复的数据，其中一个在首选值中，我想保留这个，不管它是第一个还是最后一个。如果两个都在首选值中，那就按照通常的.drop_duplicates规则来处理。如果两个都不在首选值中，那也还是按照通常的规则来处理。

我尝试过用掩码来解决这个问题，但总是搞不对。我觉得这可能不是正确的方法。以下是我尝试过的代码。

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)
""" Produces:
  id  val tag
0  A   10   X
2  A   11   X
4  B   20   Z
6  C   30   X
7  D   40   Z

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

Answer 1

我会使用 isin 和 idxmax 来处理每个 "id/val"，这样就可以模拟一个 drop_duplicates(how="not preferred") 的效果：

def pref_dup(g):
    return g.loc[[g.isin(preferred_tags).idxmax()]]

out = (
    df
     .groupby(["id", "val"])["tag"]
     .apply(pref_dup, include_groups=False)
     .reset_index(level=[0, 1])
)

输出结果：

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

Answer 2

这可能不是你想要的最优雅的方法，但这是一种可行的做法。

def conditional_remove_duplicates(df, preferred_tags):

   df = pd.DataFrame(list(pd.Series(df.itertuples(index=False)).drop_duplicates(keep='first')))
   duplicates_mask = df.duplicated(subset=['id', 'val'], keep=False)
   preferred_mask = df['tag'].isin(preferred_tags)
   mask = (~duplicates_mask | preferred_mask ).rename('mask')
   df = df[mask].drop_duplicates(subset=['id', 'val'], keep='first')

return df

我对Python还很陌生，所以任何评论或建议都非常欢迎。

Answer 3

你应该创建一个列，用来表示你想给重复项的优先级，然后根据这个优先级进行排序。这里的关键是使用 mergesort 排序，因为它是稳定的，不会改变当前相同项的顺序。

接着，你就可以去掉重复的项了。

df['rank'] = df['tag'].isin(preferred_tags)
df = df.sort_values(by='rank', ascending=False, kind='mergesort')
df.drop_duplicates(['id', 'val'])

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