Pandas groupby with None

Question

For a dataframe df

df = pd.DataFrame({'id': ['1', '1', None, None, '1', '2', '2', '3', None, '4'], 
               'last_name': ['Clinton', 'Clinton', 'Clinton','Clinton', None, 'Bush', 'Bush', None, 'Obama', 'Obama'],
               'first_name': ['Bill', 'William', 'Bill', 'William', None, 'Georg W.', 'Georg', None, 'Barack', 'Barack']})

df['id'] = df['id'].astype('category')
print(df)

which gives the following table

    id last_name first_name
0    1   Clinton       Bill
1    1   Clinton    William
2  NaN   Clinton       Bill
3  NaN   Clinton    William
4    1       NaN       None
5    2      Bush   Georg W.
6    2      Bush      Georg
7    3       NaN       None
8  NaN     Obama     Barack
9    4     Obama     Barack

I want to group by the id and last_name, drop duplicates, and remove None iff there is more than one entry. So the output should be like

              first_name
id  last_name           
1   Clinton       Bill
    Clinton       William
2   Bush          Georg W.
    Bush          Georg
3   None          None
4   Obama         Barack

One of my problems is that groupby does not work, because it excludes the None / NaN values.

Any elegant ideas?

Answer 1

IIUC, assuming your data frame has the structure similar to the one you posted, you can use ffill() and group by it, and then dropna only if len of each group is greater than 1.

df.groupby([df.id.ffill(), df.last_name.ffill()]).apply(lambda s: s.dropna() if len(s) > 1 else s).reset_index(drop=True)

    id  last_name   first_name  id2
0   1   Clinton     Bill        1
1   1   Clinton     William     1
2   2   Bush        Georg W.    2
3   2   Bush        Georg       2
4   3   None        None        3
5   NaN Obama       Barack      3