python pandas deduplication with complex criteria

user_9463364 0
Viewed 98

I have a dataframe below:

import pandas as pd
d = {'id': [1, 2, 3, 4, 4, 6, 1, 8, 9], 'cluster': [7, 2, 3, 3, 3, 6, 7, 8, 8]}
df = pd.DataFrame(data=d)
df = df.sort_values('cluster')

I want to keep ALL the rows if there is the same cluster but different id AND keep every row from that cluster even if it is the same id since there was a different id AT LEAST once within that cluster. The code I have been using to achieve this is the following below, BUT, the only problem with this is it drops too many rows for what I am looking for.

df = (df.assign(counts=df.count(axis=1))
   .sort_values(['id', 'counts'])
   .drop_duplicates(['id','cluster'], keep='last')
   .drop('counts', axis=1))

The output dataframe I am expecting that the code above does not do would drop rows at dataframe index 1, 5, 0, and 6 but leave dataframe indexes 2, 3, 4, 7, and 8. Essentially resulting in what the code below produces:

df = df.loc[[2, 3, 4, 7, 8]]

I have looked at many deduplication pandas posts on stack overflow but have yet to find this scenario. Any help would be greatly appreciated.

pandas dataframe python
1 Answers

I think we can do this with a single boolean. using .groupby().nunique()

con1 = df.groupby('cluster')['id'].nunique() > 1

#of these we only want the True indexes.

cluster
2    False
3     True
6    False
7    False
8     True


df.loc[(df['cluster'].isin(con1[con1].index))]

   id  cluster
2   3        3
3   4        3
4   4        3
7   8        8
8   9        8
user_93751020
Related