如何优化基于经纬度区间分配Segment的DataFrame函数性能

我编写了一个函数，用于为DataFrame的每行分配对应的纬度和经度区间类别，但该函数运行速度极慢。请问如何提升其性能？

原代码如下：

def assign_segment(use_df: pd.DataFrame, 
                   lat_categories: pd.core.indexes.interval.IntervalIndex, 
                   lng_categories: pd.core.indexes.interval.IntervalIndex) -> pd.DataFrame:
    """
    根据"use_tb"的纬度和经度列分配Segment。

    参数
    ----------
    use_df : pd.DataFrame
        待处理的DataFrame。
    lat_categories : pd.core.indexes.interval.IntervalIndex
        纬度区间类别。
        (示例) IntervalIndex([(35.809, 35.816], (35.816, 35.824], 
                             (35.824, 35.832], (35.832, 35.84], (35.84, 35.848]])
    lng_categories : pd.core.indexes.interval.IntervalIndex
        经度区间类别。
        (示例) IntervalIndex([(128.668, 128.685], (128.685, 128.703], 
                             (128.703, 128.72], (128.72, 128.737]])

    返回值
    -------
    use_df : pd.DataFrame
        已分配Segment的"use_df"。
    """
    segment = []

    # 遍历每行并根据纬度和经度获取对应的Segment
    for idx, row in use_df.iterrows():
        use_lat = row['use_lat']
        use_lng = row['use_lng']

        for lat_idx, lat_category in enumerate(lat_categories):
            if use_lat in lat_category:
                lat_segment = lat_idx + 1
                break
        for lng_idx, lng_category in enumerate(lng_categories):
            if use_lng in lng_category:
                lng_segment = lng_idx + 1
                break

        num_lng_grid = len(lat_categories)      # 经度网格数量
        lng_num_digits = len(str(num_lng_grid)) # 经度网格数量的位数
        segment.append((lat_segment*10**lng_num_digits)+lng_segment)
        
    # 创建segment列并赋值为函数中生成的segment列表
    use_df['segment'] = segment

    return use_df

原代码性能差的核心原因是逐行遍历（iterrows）+ 嵌套循环查找区间，完全没有利用Pandas的矢量化运算能力。以下是两种高效优化方案：

方案1：使用pd.cut矢量化映射区间 #

pd.cut是Pandas内置的区间映射工具，底层基于C实现，能一次性处理整列数据，效率比逐行循环高几个数量级。

优化后的代码：

import pandas as pd

def assign_segment_optimized(use_df: pd.DataFrame, 
                             lat_categories: pd.core.indexes.interval.IntervalIndex, 
                             lng_categories: pd.core.indexes.interval.IntervalIndex) -> pd.DataFrame:
    # 矢量化获取纬度区间对应的segment（+1保持和原逻辑一致）
    lat_segment = pd.cut(use_df['use_lat'], bins=lat_categories, labels=False) + 1
    # 矢量化获取经度区间对应的segment（+1保持和原逻辑一致）
    lng_segment = pd.cut(use_df['use_lng'], bins=lng_categories, labels=False) + 1
    
    # 计算最终segment值，逻辑与原代码完全一致
    num_lng_grid = len(lat_categories)
    lng_num_digits = len(str(num_lng_grid))
    use_df['segment'] = lat_segment * (10 ** lng_num_digits) + lng_segment
    
    return use_df

关键说明： #

• pd.cut的labels=False参数会返回每个值对应的区间索引（从0开始），加1后和原代码的lat_idx +1逻辑匹配；
• 整列操作避免了逐行循环，大幅提升处理速度，尤其是当DataFrame行数较多时效果更明显。

方案2：区间等宽时用数学计算直接推导（性能极致优化） #

如果你的经纬度区间是连续且等宽的，可以跳过区间查找，直接通过数学计算得到segment，速度比pd.cut更快。

示例代码（假设纬度区间等宽）：

def assign_segment_extreme_optimized(use_df: pd.DataFrame, 
                                     lat_categories: pd.core.indexes.interval.IntervalIndex, 
                                     lng_categories: pd.core.indexes.interval.IntervalIndex) -> pd.DataFrame:
    # 计算纬度区间的起始值和步长
    lat_start = lat_categories.left.min()
    lat_step = lat_categories[0].right - lat_categories[0].left
    # 直接计算纬度segment
    lat_segment = ((use_df['use_lat'] - lat_start) // lat_step).astype(int) + 1
    
    # 同理计算经度segment
    lng_start = lng_categories.left.min()
    lng_step = lng_categories[0].right - lng_categories[0].left
    lng_segment = ((use_df['use_lng'] - lng_start) // lng_step).astype(int) + 1
    
    # 计算最终segment值
    num_lng_grid = len(lat_categories)
    lng_num_digits = len(str(num_lng_grid))
    use_df['segment'] = lat_segment * (10 ** lng_num_digits) + lng_segment
    
    return use_df

注意事项： #

• 该方案仅适用于区间连续且等宽的场景，如果区间不等宽，计算结果会出错；
• 若区间是左开右闭（如示例中的(a, b]），需确保所有数据都落在区间内，避免出现索引偏移。

额外优化提示 #

• 原代码中num_lng_grid和lng_num_digits在循环内重复计算，这两个值是固定的，原代码中可以提前放到循环外计算，减少冗余运算；
• 日常处理Pandas数据时，尽量优先使用矢量化操作（如pd.cut、直接列运算），避免iterrows/itertuples这类逐行遍历方式。

内容的提问来源于stack exchange，提问作者Yun Tae Hwang