如何加速Pandas中单元格坐标计算的循环数据提取代码？

问题描述 #

face_Data是包含Cell_0、Cell_1（均为单元格ID）及X、Y坐标列的DataFrame，目标是通过它计算每个单元格的坐标。cell_Data的ID列为连续序列[1,2,…N]，当前通过循环遍历每个ID查询对应面数据并求均值，但face_Data规模通常达10^5或106量级，循环耗时极长，急需非循环的加速方法。

现有代码 #

X = []
Y = []
for i in cell_Data['ID']:
    A = face_Data.query(f"Cell_0 == {i} or Cell_1 == {i}")
    X.append(np.average(A['X']))
    Y.append(np.average(A['Y']))

cell_Data['X'] = X
cell_Data['Y'] = Y
print("Cell Coordinates Obtained")

示例数据 #

face_data = pd.DataFrame({
             'ID': [1, 2, 3, 4, 5, 6],
             'X': [0, 0.1, 0.2, 0, 0.1, 0.2],
             'Y': [1, 1, 1, 2, 2, 2],
             'Cell_0': [1, 2, 3, 5, 6, 7],
             'Cell_1': [2, 3, 4, 6, 7, 8]
            })
cell_data = pd.DataFrame([1, 2, 3, 4, 5, 6, 7, 8], columns=['ID'])

注：数据基于二维笛卡尔网格，括号内为cell_ID，竖线下为face_ID。

优化方案 #

核心思路：重塑数据实现向量化计算 #

把Cell_0和Cell_1列拆分成多行，让每个单元格ID对应其关联的所有面坐标，之后直接按ID分组求均值，完全规避循环操作。

代码实现 #

import pandas as pd

# 重塑数据：将Cell_0、Cell_1转为长格式，每个单元格ID对应一条坐标记录
melted = pd.melt(
    face_Data,
    id_vars=['X', 'Y'],
    value_vars=['Cell_0', 'Cell_1'],
    var_name='Cell_Type',
    value_name='Cell_ID'
)

# 按单元格ID分组，计算X、Y坐标的均值
cell_coords = melted.groupby('Cell_ID')[['X', 'Y']].mean().reset_index()

# 将计算结果合并到cell_Data中
cell_Data = cell_Data.merge(cell_coords, left_on='ID', right_on='Cell_ID', how='left').drop('Cell_ID', axis=1)

print("Cell Coordinates Obtained")

性能优势 #

这种向量化操作依托pandas的C级内部优化，在百万级数据量下，速度比原循环方法快几十到上百倍，彻底解决性能瓶颈。

补充说明 #

如果cell_Data中存在face_Data未覆盖的ID，how='left'会保留这些行，对应的X、Y值会设为NaN，可根据实际需求调整合并方式。

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