Pandas Groupby применить список из столбца на основе двоичного столбца

Question

У меня есть фрейм данных:

id  to    from    flag
1    a     x        1
1    a     y        0
2    c     z        1
2    c     m        1
2    b     v        0
2    b     p        0

, и я хочу сгруппировать (['id', 'to']) и вернуть список элементов из которых имеют только флаг 1.Если ни один элемент не имеет флага 1, то результатом будет «Нет».Желаемый результат должен быть:

id  to  from 
 1   a  ['x']  
 2   c  ['z','m'] 
 2   b  None

Я могу сделать это с apply, т.е.

out_df = df.groupby(['id', 'to'])['from'].apply(
       lambda x: match_to_list(x['from'], x['flag'])).reset_index()

, где:

def match_to_list(to, flag):
    matches = list(to.iloc[flag.nonzero()[0]])
    if len(matches) == 0:
        return 'None'
    else:
        matches

, но это занимает слишком много времени, и ядумаю, что должен быть лучший способ, который я пропускаю.

Буду очень признателен за любую помощь / понимание!ТИА

Answer 1

IIUC 1-ый создайте индекс с помощью MultiIndex, затем мы сделаем groupby с agg

idx=pd.MultiIndex.from_tuples(list(map(tuple,df[['id','to']].drop_duplicates().values.tolist())))
yourdf=df.loc[df.flag==1].groupby(['id','to'])['from'].agg(list).reindex(idx).reset_index()
yourdf
Out[13]: 
   level_0 level_1    from
0        1       a     [x]
1        2       c  [z, m]
2        2       b     NaN

Или просто с помощью применения, менее эффективно, но более читабельно

df.groupby(['id','to']).apply(lambda x : x['from'][x['flag']==1].tolist() if (x['flag']==1).any() else None).reset_index()
Out[17]: 
   id to       0
0   1  a     [x]
1   2  b    None
2   2  c  [z, m]