如何加速Python函数调用：DataFrame邮政编码批量处理优化

问题背景 #

我有一个数据集，核心字段为postalCode（整数类型），需要生成新字段mainPostCode并添加到DataFrame中。mainPostCode是对应区域的主邮局编码，示例如下：

生成规则 #

提取postalCode的前三位数字，在包含约40000个元素的zipcode列表（示例：zipcode = [1001,1002,...,99999]）中，找出以该三位数字开头且小于当前postalCode的最小值；若不存在符合条件的编码，则返回当前postalCode。

当前慢代码 #

当前使用的函数及调用方式处理数据集预计耗时72小时，性能瓶颈明显：

def findMainPostOffice(num):
    ''' takes zip and returns nearest available main zip in list 'zipcode' '''
    start = int(str(num // 100) + '00')
    m = min([i for i in zipcode if i > start and i < num], default=num)
    return m

调用代码：

df['mainPostCode'] = df.postalCode.apply(findMainPostOffice)

加速方案 #

1. 预分组+二分查找（大幅提速） #

原函数每次查询都遍历整个4万条的zipcode列表，时间复杂度为O(NM)（N为数据集行数，M为zipcode长度）。通过预分组+二分查找，将时间复杂度降至O(Nlog K)（K为每组平均元素数，约44）。

步骤1：预处理zipcode

import bisect
from collections import defaultdict

# 先对zipcode整体排序
zipcode_sorted = sorted(zipcode)

# 按前三位数字分组，每组内保持排序
zip_groups = defaultdict(list)
for z in zipcode_sorted:
    prefix = z // 100  # 提取前三位（如12345//100=123）
    zip_groups[prefix].append(z)

步骤2：改写查找函数

def findMainPostOffice_fast(num):
    prefix = num // 100
    group = zip_groups.get(prefix, [])
    start = prefix * 100
    
    if not group:
        return num
    
    # 用二分查找定位第一个大于start的元素位置
    idx = bisect.bisect_right(group, start)
    # 若该元素小于当前num，即为目标最小值；否则返回原num
    if idx < len(group) and group[idx] < num:
        return group[idx]
    else:
        return num

调用方式

df['mainPostCode'] = df.postalCode.apply(findMainPostOffice_fast)

2. NumPy优化二分查找（进一步提速） #

如果数据集规模极大，可将分组转成NumPy数组，利用更高效的底层实现加速：

import numpy as np

# 将分组转成NumPy数组
zip_groups_np = {k: np.array(v) for k, v in zip_groups.items()}

def findMainPostOffice_np(num):
    prefix = num // 100
    arr = zip_groups_np.get(prefix, np.array([]))
    start = prefix * 100
    
    if arr.size == 0:
        return num
    
    # NumPy二分查找
    idx = np.searchsorted(arr, start, side='right')
    if idx < arr.size and arr[idx] < num:
        return arr[idx].item()
    else:
        return num

调用方式

df['mainPostCode'] = df.postalCode.apply(findMainPostOffice_np)

3. 全向量化操作（适合超大规模数据集） #

避免apply循环，用Pandas分组+向量化逻辑处理，彻底消除Python循环开销：

import pandas as pd

# 原数据添加前缀和起始值列
df['prefix'] = df['postalCode'] // 100
df['start'] = df['prefix'] * 100

# 定义分组处理函数
def get_main_zip(group):
    postal_codes = group['postalCode'].values
    start_val = group['start'].iloc[0]
    zips = zip_groups[group.name]
    
    results = []
    for pc in postal_codes:
        idx = bisect.bisect_right(zips, start_val)
        if idx < len(zips) and zips[idx] < pc:
            results.append(zips[idx])
        else:
            results.append(pc)
    return results

# 分组计算并赋值
df['mainPostCode'] = df.groupby('prefix').apply(get_main_zip).explode()

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