python - 如何加快python函数调用速度

Question

我有一个数据集，为简单起见，我将仅指出一个主要特征 - 邮政编码。 我需要通过函数调用获得另一项功能(该地区的主要邮局)并添加到数据框(示例)中。

两者都是整数。

<表类=“s-表”> <标题> 邮政编码 主邮政编码 <正文> 12345 12301 23456 23407 34567 34504

关于函数的一些说明:它获取邮政编码的前 3 位数字，然后从所有邮政编码列表中获取从这 3 位数字开始的最小值。

您不会总是在此列表中找到看起来像 XXX01 的值，它可以是 XXX05 或 XXX07 或 XXX(任何其他)。我们假设它可以是任何数字。

邮政编码列表如下所示(大约 40K 元素):

zipcode = [1001,1002,...,99999]

我的函数如下所示:

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)

问题是这个函数需要很长时间。在我的数据集上，大约需要 72 小时。 您能帮我加快速度吗？

Answer 1

IIUC，可以使用groupby求最小值(主要邮政编码)

df['mainPostCode'] = (df.groupby(df['postalCode'].astype(str).str.zfill(5).str[:2])
                        .transform('min'))
print(df)

# Output
       postalCode  mainPostCode
0           23041         23003
1           48558         48000
2           52895         52000
3           39817         39000
4           40427         40000
...           ...           ...
39995       81184         81000
39996        7125          7001
39997       22773         22003
39998       88802         88002
39999       58510         58000

[40000 rows x 2 columns]

输入:

import pandas as pd
import numpy as np

np.random.seed(2023)
df = pd.DataFrame({'postalCode': np.random.randint(1000, 100000, 40000)})