Как добавить два списка категорий в словарь в python

Question

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

df =
          name     age     character
0          A        10       fire
1          A        15       water
2          A        20       earth
3          A        25       air
4          B        10       fire
5          B        7        air

Я хочу преобразовать этот фрейм данных в словарь, чтобы получился вывод,

dic = {'A': [[10, 15, 20, 25], ['fire', 'water', 'earth', 'air']],
       'B': [[10, 7], ['fire', 'air']] }

То, что я пробовал,

from collections import defaultdict
dic = defaultdict(list)
for i in range(len(df)):
    dic[df.loc['name', i]].append(df.loc['age', i])
    dic[df.loc['name', i]].append(df.loc['character', i]) # this is wrong. It appends to existing list.

Если я объявляю dic = defaultdict([[], []]), выдается ошибка, что первый аргумент defaultdict должен быть вызван или Нет. Как я могу улучшить этот словарь?

Answer 1

Вы можете использовать комбинацию pivot_table и to_dict

dic = (df.pivot_table(columns='name', values=['age','character'], aggfunc=list)
         .to_dict('l'))

Out[107]:
{'A': [[10, 15, 20, 25], ['fire', 'water', 'earth', 'air']],
 'B': [[10, 7], ['fire', 'air']]}

Если у вашего фрейма данных точно 3 столбца name, age, character, вы можете просто игнорировать values= параметр

dic = df.pivot_table(columns='name', aggfunc=list).to_dict('l')

---

Как вы сказали в комментарии, для удаления пробелов необходимо предварительно обработать df с помощью str.strip перед вызовом pivot_table следующим образом

df.update(df.select_dtypes('object').apply(lambda x: x.str.strip()))
dic = df.pivot_table(columns='name', aggfunc=list).to_dict('l')

Answer 2

Вот решение, которое возвращает np.array, что достаточно похоже на list:

{k: d[['age','character']].T.to_numpy() for k,d in df.groupby('name')}

Вывод:

{'A': array([[10, 15, 20, 25],
        ['fire', 'water', 'earth', 'air']], dtype=object), 
'B': array([[10, 7],
        ['fire', 'air']], dtype=object)}