Pandas dataframe groupby make a list or array of a column

Question

import pandas as pd
import numpy as np

df = {'a': ['aa', 'aa', 'aa', 'aaa', 'aaa'], 
      'b':['bb', 'bb', 'bb', 'bbb', 'bbb'], 
      'c':[10,20,30,100,200]}

df = pd.DataFrame(data=df)

my_dict=df.groupby(['a', 'b'])['c'].apply(np.hstack).to_dict()

gives the following dictionary

>>> my_dict
{('aa', 'bb'): array([10, 20, 30]), ('aaa', 'bbb'): array([100, 200])}

Is there a faster/efficient way of doing this other than using apply?

jezrael · Accepted Answer · 2020-09-17 09:05:36Z

5

Use dictionary comprehension:

my_dict= {k:np.hstack(v) for k, v in df.groupby(['a', 'b'])['c']}
print (my_dict)
{('aa', 'bb'): array([10, 20, 30]), ('aaa', 'bbb'): array([100, 200])}

answered Sep 17, 2020 at 9:05

jezrael

868k103 gold badges1.4k silver badges1.3k bronze badges

Sign up to request clarification or add additional context in comments.

Comments

halfer · Accepted Answer · 2022-05-22 11:45:39Z

2

You could use groupby and itertuples:

my_dict = dict(df.groupby(['a','b']).agg(list).itertuples(name=None))

{('aa', 'bb'): [10, 20, 30], ('aaa', 'bbb'): [100, 200]}

Or more succinctly, as noted by Ch3steR:

df.groupby(['a','b']).agg(list).to_dict() 


{('aa', 'bb'): [10, 20, 30], ('aaa', 'bbb'): [100, 200]}

edited May 22, 2022 at 11:45

halfer

20.2k20 gold badges110 silver badges207 bronze badges

answered Sep 17, 2020 at 9:09

Umar.H

23.1k7 gold badges50 silver badges94 bronze badges

2 Comments

Ch3steR Over a year ago

df.groupby(['a','b']).agg(list).to_dict() -> {('aa', 'bb'): [10, 20, 30], ('aaa', 'bbb'): [100, 200]}

Umar.H Over a year ago

@Ch3steR much cleaner thanks, added as an answer. I guess indexes get returned as tuples when using .to_dict methods.

Collectives™ on Stack Overflow

Pandas dataframe groupby make a list or array of a column

2 Answers 2

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related