在Pandas DataFrame单元格创建NumPy数组并高效实现元素乘加运算

问题背景 #

已有如下DataFrame：

import pandas as pd
import numpy as np

d = {'col1': [[1,2,3]]}
df = pd.DataFrame(d)

通过以下代码生成了包含NumPy数组的squared和cubed列：

df['squared'] = [np.array([i**2 for i in x]) for x in df['col1']]
df['cubed'] = [np.array([i**3 for i in x]) for x in df['col1']]

得到的DataFrame如下：

col1   squared      cubed
0  [1, 2, 3] [1, 4, 9] [1, 8, 27]

需求：对squared和cubed列中的每个NumPy数组执行元素级乘法后求和，且需适配大数据量场景的高效实现。

高效实现方案 #

核心思路：利用NumPy向量化运算 #

避免Python层面的逐行循环，将列中的NumPy数组批量转换为二维矩阵，通过NumPy的向量化操作完成计算——底层由C实现的运算能大幅降低开销，适配大数据量场景。

代码实现 #

# 将列中的所有NumPy数组堆叠为二维矩阵
squared_matrix = np.vstack(df['squared'].values)
cubed_matrix = np.vstack(df['cubed'].values)

# 执行元素级乘法后按行求和，结果存入新列
df['elementwise_mult_sum'] = (squared_matrix * cubed_matrix).sum(axis=1)

执行后，elementwise_mult_sum列的结果为276（对应计算：1*1 + 4*8 + 9*27 = 1 + 32 + 243）。

性能优势 #

当数据量较大时，这种方法的性能远优于df.apply或列表推导式：

• NumPy的向量化运算跳过了Python循环的解释开销，直接调用底层优化的C代码
• 能充分利用CPU的向量处理能力，批量完成矩阵运算，效率提升显著

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