Быстрое преобразование многоколоночного фрейма данных в словарь

Question

У меня следующая проблема. У меня есть pandas фрейм данных со столбцами от A до D, где столбцы A и B являются своего рода идентификатором. Моя конечная цель - создать словарь, в котором кортеж (A, B) обозначает ключи, а значения C и D хранятся под каждым ключом как массив numpy. Я могу написать это в одну строку, если я хочу сохранить только C или D, но я изо всех сил пытаюсь заполучить их обоих. Вот что у меня есть:

   output_dict = df.groupby(['A','B'])['C'].apply(np.array).to_dict()

работает как положено, то есть данные для каждого ключа представляют собой массив dim (N, 1). Но если я попытаюсь сделать следующее:

   output_dict = df.groupby(['A','B'])['C','D'].apply(np.array).to_dict()

Я получу сообщение об ошибке:

TypeError: Series.name must be a hashable type

Как я могу включить 2-й столбец так, чтобы данные в ключе для каждого ключа представляли собой массив дим (N, 2).

Спасибо!

Answer 1

Я немного поиграл, и рядом с уже полезным ответом Герда я нашел следующее, соответствующее моим потребностям, используя лямбду.

output_dict = df.groupby(['A','B']).apply(lambda df: np.array( [ df['C'],df['D'] ] ).T).to_dict()

Сравнение времени с решением Герда в моем конкретном случае: Герд: примерно 0,055 s Это: примерно 0,035 с

Answer 2

Вы можете создать новый столбец (например, C_D), содержащий списки соответствующих значений в столбцах C и D. Выберите столбцы C и D из кадра данных и используйте метод tolist():

df['C_D'] = df[['C','D']].values.tolist()

Затем запустите строку кода для этого нового столбца:

output_dict = df.groupby(['A','B'])['C_D'].apply(np.array).to_dict()