Python Pandas создает столбец на основе функции, требующей фильтрации DataFrame

Question

У меня есть большой скрипт с несколькими функциями.В одной из этих функций я создаю фрейм данных, а затем создаю столбец с применением отдельной функции.

Функция для создания фрейма данных на высоком уровне:

def data(file):
  df = pd.DataFrame('A': [1,2,3,4], 'B':[5,5,6,6]
  df['C'] = df['B'].apply(func)

Функция 'func'по сути, предполагается фильтровать кадр данных по столбцу B и возвращать список значений в столбце 'A'

def func(x):
  df2 = df[df['B']==x]
  names = df2['A']
  return names

К сожалению, я не могу использовать глобальный вызов для извлечения df в функцию, поэтому я запутался, каквыполнить этот запрос.Идеальный выход должен быть таким:

A B C
1 5 [1,2]
2 5 [1,2]
3 6 [3,4]
4 6 [3,4]

Answer 1

Использование map после groupby.apply (PS: не рекомендуется использовать список в столбце, что усложнит настройку)

df['C']=df.B.map(df.groupby('B').A.apply(list))
df
Out[872]: 
   A  B       C
0  1  5  [1, 2]
1  2  5  [1, 2]
2  3  6  [3, 4]
3  4  6  [3, 4]