优化大型Pandas DataFrame中4个XYZ坐标的定位效率

• 数据规模：50万条3D网格点数据存储在Pandas DataFrame中
• 查找规则：需按以下规则定位4个特定坐标点：
  1. 第一个点：指定半径、0度位置
  2. 第二个点：与第一个点Y、Z坐标一致，取另一可选X值（仅两个X可选）
  3. 第三个点：与第一个点X值一致，Y、Z为其相反数（对应180度位置）
  4. 第四个点：与第三个点Y、Z坐标一致，取另一可选X值
• 数据特征：来自3个CSV文件，共享Node Number、X/Y/Z坐标字段，仅力的方向列不同；点分布在不同半径的圆周上，间隔11.25度，Y代表长度、Z代表高度、X代表深度
• 当前问题：使用多层循环实现单半径查找耗时11分钟，存在严重性能瓶颈，寻求高效优化方案

当前实现代码 #

# 查找4个点组合的函数
# df_points: 存储所有节点位置的DataFrame
# points_arr: 0-180度范围内的所有点数组
def find_point_combinations(df_points, points_arr):
    
    # 存储点组合的数组
    voltage = []

    for n in range(0, len(points_arr)):
        point1 = points_arr[n]
        point3 = point1

        # 查找point3
        for pt in range(0, len(df_points)):
            if df_points['Y Location (mm)'].loc[pt] == -1 * point1[2] and df_points['Z Location (mm)'].loc[pt] == -1 * point1[3]:
                if df_points['X Location (mm)'].loc[pt] == point1[1]:
                    point3 = df_points.loc[pt]
                    break  

        # 查找对面的point4
        for pt in range(0, len(df_points)):
            if df_points['Y Location (mm)'].loc[pt] == point3[2] and df_points['Z Location (mm)'].loc[pt] == point3[3]:
                if df_points['X Location (mm)'].loc[pt] != point3[1]:
                    point4 = df_points.loc[pt]
                    break

        # 查找point2
        for pt in range(0, len(df_points)):
            if df_points['Y Location (mm)'].loc[pt] == point1[2] and df_points['Z Location (mm)'].loc[pt] == point1[3]:
                if df_points['X Location (mm)'].loc[pt] != point1[1]:
                    point2 = df_points.loc[pt]
                    break

        voltage.append([point1, point2, point3, point4])
    return voltage

高效优化方案 #

核心思路：用Pandas索引匹配和向量化操作替代嵌套循环，将时间复杂度从O(N*M)降至O(N)级别。

1. 预处理：建立复合索引与映射表 #

提前构建索引和映射关系，避免重复遍历数据：

import pandas as pd

# 建立(X,Y,Z)复合索引，支持快速定位单条数据
df_points = df_points.set_index(['X Location (mm)', 'Y Location (mm)', 'Z Location (mm)'])

# 建立(Y,Z)到对应X值列表的映射（每个(Y,Z)仅两个X值）
yz_to_x = df_points.reset_index().groupby(['Y Location (mm)', 'Z Location (mm)'])['X Location (mm)'].apply(list).to_dict()

2. 批量查找点组合 #

利用索引直接定位目标点，彻底消除嵌套循环：

def find_point_combinations_optimized(df_points, points_arr, yz_to_x):
    voltage = []
    
    for point1 in points_arr:
        # 解析point1的坐标（假设结构为[Node Number, X, Y, Z, ...]）
        x1, y1, z1 = point1[1], point1[2], point1[3]
        
        # 查找point3：X=x1，Y=-y1，Z=-z1
        try:
            point3 = df_points.loc[(x1, -y1, -z1)]
        except KeyError:
            point3 = None  # 可根据业务需求处理未找到的情况
        
        # 查找point2：Y=y1、Z=z1，取另一个X值
        x_candidates = yz_to_x.get((y1, z1), [])
        x2 = next((x for x in x_candidates if x != x1), None)
        point2 = df_points.loc[(x2, y1, z1)] if x2 is not None else None
        
        # 查找point4：Y=-y1、Z=-z1，取另一个X值
        point4 = None
        if point3 is not None:
            x3 = point3.name[0]
            x_candidates_p4 = yz_to_x.get((-y1, -z1), [])
            x4 = next((x for x in x_candidates_p4 if x != x3), None)
            if x4 is not None:
                point4 = df_points.loc[(x4, -y1, -z1)]
        
        voltage.append([point1, point2, point3, point4])
    
    return voltage

3. 额外优化：合并CSV与数据类型压缩 #

• 合并3个CSV文件，避免重复加载相同坐标数据：

dfs = [pd.read_csv(f'file{i}.csv') for i in range(1,4)]
df_combined = pd.merge(dfs[0], dfs[1], on=['Node Number', 'X Location (mm)', 'Y Location (mm)', 'Z Location (mm)'], how='outer')
df_combined = pd.merge(df_combined, dfs[2], on=['Node Number', 'X Location (mm)', 'Y Location (mm)', 'Z Location (mm)'], how='outer')

• 压缩数值列数据类型，减少内存占用并提升操作速度：

for col in ['X Location (mm)', 'Y Location (mm)', 'Z Location (mm)']:
    df_combined[col] = pd.to_numeric(df_combined[col], downcast='float')
df_combined['Node Number'] = pd.to_numeric(df_combined['Node Number'], downcast='integer')

示例输出 #

Node Number                 2.708330e+05
X Location (mm)             1.168100e+02
Y Location (mm)             2.060000e-04
Z Location (mm)             2.958000e+02
X Elastic Strain (mm/mm)             NaN
Y Elastic Strain (mm/mm)    2.420000e-06
Z Elastic Strain (mm/mm)   -5.970000e-07
Error                       4.200000e+00
Angle (deg)                 1.800000e+02
Radius                      2.958000e+02

Name: (116.81, 0.000206, 295.8), dtype: float64
Node Number                 90690.000000
X Location (mm)               124.010000
Y Location (mm)               -58.527000
Z Location (mm)               294.240000
X Elastic Strain (mm/mm)             NaN
Y Elastic Strain (mm/mm)        0.000005
Z Elastic Strain (mm/mm)        0.000001
Error                         597.115221
Angle (deg)                   180.000000
Radius                        295.800000

Name: (124.01, -58.527, 294.24), dtype: float64
Node Number                 2.708330e+05
X Location (mm)             1.168100e+02
Y Location (mm)            -2.060000e-04
Z Location (mm)            -2.958000e+02
X Elastic Strain (mm/mm)             NaN
Y Elastic Strain (mm/mm)    2.420000e-06
Z Elastic Strain (mm/mm)   -5.970000e-07
Error                       4.200000e+00
Angle (deg)                 1.800000e+02
Radius                      2.958000e+02

Name: (116.81, -0.000206, -295.8), dtype: float64
Node Number                 90509.000000
X Location (mm)               124.010000
Y Location (mm)                58.527000
Z Location (mm)              -294.240000
X Elastic Strain (mm/mm)             NaN
Y Elastic Strain (mm/mm)        0.000003
Z Elastic Strain (mm/mm)       -0.000001
Error                          58.809755
Angle (deg)                   180.000000
Radius                        295.800000

Name: (124.01, 58.527, -294.24), dtype: float64

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