Grouping duplicates, allowing NaN to be equal to any value

Viewed 602

I want to aggregate duplicates rows in a DataFrame based on multiple columns, but with a non-regular definition of what 'duplicate' means.

Imagine a DataFrame with two columns, A and B:


In [45]: df = pd.DataFrame([[0, 0], [pd.NA, 0], [0, 1], [pd.NA, 1], [1, 1]], columns=['A', 'B']).convert_dtypes()

In [63]: df
Out[63]: 
      A  B
0     0  0
1  <NA>  0
2     0  1
3  <NA>  1
4     1  1

Basically, a missing value should be allowed to equal to any value, so long as it does not need 'be equal' to multiple values at once.

Thus, the first and second row should be declared duplicates, and aggregate together. Row 0 and 1 are duplicates because they share the same value in column B, and because there is only 1 unique value (not counting missing values) in column A.

This imposes a problem with the last three rows, since the missing value could fill in for either 0 or 1. In this case, I don't want the rows aggregated together; neither row 2 & 3 should be duplicates, nor row 3 & 4. Even though row 2, 3 and 4 share the same value in column B, none of them are duplicates of another because there are 2 unique values (not counting missing value) in column A.

The ultimate goal is aggregate duplicate rows together, like with DataFrame.groupby().aggregate(), but if you know a solution for just tagging rows as duplicate, returning a boolean Series like DataFrame.duplicated, then that gets me some of the way there.

For my problem in particular, there will be arbitrary many columns, but only 1 column with missing values. You may assume that there exists a unique fill-value to fill in the missing values in that column to avoid those rows being dropped in a groupby.

1 Answers

one way could be to groupby all your other columns (here only B) and transform with False if nunique in the group is 1 (nan are not counted) and True otherwise. Then you can use this mask or notna to keep the rows you want:

print (df.loc[ df.groupby('B')['A'].transform(lambda x: False if x.nunique()==1 else True)
                |df['A'].notna(), :])
      A  B
0     0  0
2     0  1
3  <NA>  1
4     1  1

EDIT: create the group if only one value other than np.nan could be done with transform and replace nan with the unique value:

df['gr'] = df.groupby('B')['A'].transform(lambda x: x.max() if x.nunique()==1 else x)

then you can perform some groupby on B and this column gr, not sure of your expected output

Related