如何对Pandas DataFrame列内每个单元格的数组元素求和

问题说明 #

当DataFrame某一列存储长度不一致的数值列表时，直接将列转为numpy数组按行求和会触发维度错误，需求为对该列每个单元格内的列表做元素求和，最终每个单元格值为对应和组成的单元素列表。
示例测试数据：

import pandas as pd

d = {"squares":[[1, 1, 1], [1], [1, 1], [1, 1, 1, 1]]}
df = pd.DataFrame(d)

预期输出结果：

[3]
[1]
[2]
[4]

实现方法 #

方法1：apply映射 #

直接对列的每个元素调用内置sum函数，再包装为要求的单元素列表即可，天然兼容不等长列表场景：

df["squares"] = df["squares"].apply(lambda x: [sum(x)])

执行后打印df["squares"]即可得到匹配预期的结果。

之前numpy方案报错原因：不等长列表转numpy时不会生成二维数值数组，只会生成存储列表对象的一维object数组，指定axis=1求和时自然会触发轴维度不匹配的错误。

方法2：列表推导式 #

如果数据量较大，列表推导式的执行效率比apply方法更高，写法更简洁：

df["squares"] = [[sum(item)] for item in df["squares"]]

如果不需要保留单元素列表格式、直接需要整数类型的求和结果，去掉外层包装的方括号即可：

# 直接得到整数类型的和序列：3、1、2、4
df["squares"] = df["squares"].apply(sum)

内容的提问来源于stack exchange，提问作者ftivoaj