creating dictionary from dataframe

Question

I have below dataframe

a = [1,2,2,3,5]
b = [10,20,30,40,50]
df = pd.DataFrame(list(zip(a,b)), columns=['a','b'])

I want to generate dictionary like below

{1:10, 2:[20,30], 3:40, 5:50}

Answer 1

Use groupby aggregate + Series.to_dict:

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

{1: [10], 2: [20, 30], 3: [40], 5: [50]}

Or conditionally aggregate into a list based on number of elements in the series to get exact output:

d = df.groupby('a')['b'].agg(lambda s: list(s) if len(s) > 1 else s).to_dict()

{1: 10, 2: [20, 30], 3: 40, 5: 50}