从数据框创建字典

Question

我有以下数据框

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

我想生成如下字典

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

Answer 1

使用 groupby aggregate + Series.to_dict：

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

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

或者根据系列中的元素数量有条件地聚合成一个 list 以获得准确的输出：

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}