Python Pandas - Как объединить два кадра данных, которые имеют два разных столбца, в два столбца списка

Question

У меня есть два кадра данных.

df1:

key  column_a 
AA   AAB      
AA   AAC      
AB   ABA      
AB   ABK      

df2:

key  column_b
AA   AAA
AB   ABA
AB   ABK
AC   ACC

Я хотел объединить их, но column_a и column_b стали бы списком.

key  column_a       column_b
AA   ['AAB','AAC']  ['AAA']
AB   ['ABA','ABK']  ['ABA','ABK']
AC   NaN            ['ACC']

Answer 1

При преобразовании в списки вы теряете векторизованную функциональность, которая связана с использованием массивов NumPy, хранящихся в смежных блоках памяти.

Решение с concat, groupby и агрегирующая лямбда-функция - удалите NaN s и создайте list s:

df = pd.concat([df1, df2]).groupby('key').agg(lambda x: x.dropna().tolist())
print (df)
       column_a    column_b
key                        
AA   [AAB, AAC]       [AAA]
AB   [ABA, ABK]  [ABA, ABK]
AC           []       [ACC]

Последнее, если необходимо заменить пустые списки на NaN s, использовать DataFrame.where:

df = df.where(df.astype(bool))
print (df)
       column_a    column_b
key                        
AA   [AAB, AAC]       [AAA]
AB   [ABA, ABK]  [ABA, ABK]
AC          NaN       [ACC]