CODEWITHSUNDEEP
pandasdataframepandas-groupby
charmander123
charmander123

Reputation: 367

Pandas Groupby: return dict of rows

I would like to group my dataframe by one of the columns and then return a dictionary that has a list of all of the rows per column value. Is there a fast Pandas idiom for doing this?

Example:

test = pd.DataFrame({
    'id': ['alice', 'bob', 'bob', 'charlie'],
    'transaction_date': ['2020-01-01', '2020-01-01', '2020-01-02', '2020-01-02'],
    'amount': [50.0, 10.0, 12.0, 13.0]
})

Desired output:

result = {
  'alice': [Series(transaction_date='2020-01-01', amount=50.0)],
  'bob': [Series(transaction_date='2020-01-01', amount=10.0), Series(transaction_date='2020-01-02', amount=12.0)],
  'charlie': [Series(transaction_date='2020-01-02', amount=53.0)],
}

The following approaches do NOT work:

test.groupby('id').agg(list)

test.groupby('id').agg(list).to_dict():
{'amount': {'charlie': [13.0], 'bob': [10.0, 12.0], 'alice': [50.0]}, 'transaction_date': {'charlie': ['2020-01-02'], 'bob': ['2020-01-01', '2020-01-02'], 'alice': ['2020-01-01']}}

test.groupby('id').apply(list).to_dict():
{'charlie': ['amount', 'id', 'transaction_date'], 'bob': ['amount', 'id', 'transaction_date'], 'alice': ['amount', 'id', 'transaction_date']}

Upvotes: 0

Views: 249

Answers (1)

E. Zeytinci
E. Zeytinci

Reputation: 2643

Use itertuples and zip,

import pandas as pd

test = pd.DataFrame({
    'id': ['alice', 'bob', 'bob', 'charlie'],
    'transaction_date': ['2020-01-01', '2020-01-01', '2020-01-02', '2020-01-02'],
    'amount': [50.0, 10.0, 12.0, 13.0]
})

columns = ['transaction_date', 'amount']

grouped = (test
           .groupby('id')[columns]
           .apply(lambda x: list(x.itertuples(name='Series', index=False))))


print(dict(zip(grouped.index, grouped.values)))

{
    'alice': [Series(transaction_date='2020-01-01', amount=50.0)], 
    'bob': [
        Series(transaction_date='2020-01-01', amount=10.0), 
        Series(transaction_date='2020-01-02', amount=12.0)
    ], 
    'charlie': [Series(transaction_date='2020-01-02', amount=13.0)]
}

Upvotes: 1

Related Questions