Pandas按多个列表分组并聚合

Question

我有一个数据框 df，其中有两列。我想按一列进行分组，并聚合属于同一组的列表，通过获取“平均列表”（所有列表元素的平均值），如下所示：

column_a, column_b
1,         [1,2,3]
1,         [2,5,1]
2,         [5,6,6]
3,         [2,0,1]
3,         [4,2,3]

列表始终具有相同的固定长度。期望输出应如下所示：

group, avg_list
1,     [1.5,3.5,2]
2,     [5,6,6]
3,     [3,1,2]

我知道可以在dataframe上使用groupby，然后聚合它，但我不确定该在代码的agg部分中放什么。

df.groupby('column_a').agg(?)

我会感激任何建议。

- carpediem

Answer 1

您可以通过以下方式获取每个组内列表的平均值：

s = df.groupby("column_a")["column_b"].apply(lambda x: np.array(x.tolist()).mean(axis=0))

pd.DataFrame({'group':s.index, 'avg_list':s.values})

提供：

  group avg_list
0   1   [1.5, 3.5, 2.0]
1   2   [5.0, 6.0, 6.0]
2   3   [3.0, 1.0, 2.0]

- DavideBrex

---

非常感谢你。这正是我需要的。我感激你的帮助。 - carpediem

当我读入数据时，我有一个字符串列表。结果，我收到了错误信息“TypeError: cannot perform reduce with flexible type”...因此，我将代码从np.array(x.tolist())更改为np.array(x.tolist()).astype(np.float)以解决问题。 - David Erickson

不错！我的做法是使用 literal_eval 转换列。这一行代码 df["column_b"] = df.column_b.apply(lambda s: list(ast.literal_eval(s))) 将字符串转换为列表。 - DavideBrex

Answer 2

将它们作为列表组合在一起，然后使用平均值函数来找到平均值。

df.set_index('column_a')['column_b'].groupby('column_a').apply(list).apply(lambda x: np.mean(x,0))

结果

column_a
1    [1.5, 3.5, 2.0]
2    [5.0, 6.0, 6.0]
3    [3.0, 1.0, 2.0]

- Karthik Radhakrishnan