pandas 如何加快python函数调用

p4tfgftt  于 2023-02-02  发布在  Python
关注(0)|答案(4)|浏览(94)

我有一个数据集,为了简单起见,我将只指出一个主要的功能- postalCode。我需要通过函数调用获得另一个功能(这个地区的主要邮局),并添加到dataframe(示例)。
两者都是整数。
| 邮政编码|主要邮政编码|
| - ------|- ------|
| 小行星123|小行星123|
| 小行星23456|小行星23407|
| 小行星345|小行星34504|
关于功能的一些词汇:它取postalCode的前3位数,然后从所有邮政编码的列表中取最小值,从这3位数开始。
在这个列表中,你并不总能找到类似XXX 01的值,它可以是XXX 05、XXX 07或XXX(任何其他值),让我们假设它可以是任何数字。
邮政编码列表如下所示(约40 K个元素):

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个小时。你能帮我加快这个速度吗?

yiytaume

yiytaume1#

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)})
31moq8wy

31moq8wy2#

你应该尽可能多地把计算移出函数。
对于任何前缀,我们都需要最低的主要邮政局,这样我们就可以创建一个前缀到主要邮政编码的Map,因为只有1000个可能的前缀,所以不会占用太多空间。
一种方法是为所有可能的邮政编码创建一个前缀字典,对于邮政编码[10001, 10002, 20010, 20004]的列表,我们创建一个Map:

{
    100: 10001,
    200: 20004,
}

我们不关心邮政编码1000120010,因为我们永远不会返回它们。
通过只创建一次Map,并多次使用它,我们不必在每次搜索邮政编码时检查整个列表。
下面是生成Map的代码:

zipcodes = [10000, 99999]

prefix_map = {}

for z in zipcodes:
    prefix = z // 100
    if prefix in prefix_map:
        prefix_map[prefix] = min(prefix_map[prefix], z)
    else:
        prefix_map[prefix] = z

下面是使用前缀_map的代码

def findMainPostOffice(num):
    global prefix_map
    prefix = num // 100
    if prefix in prefix_map and prefix_map[prefix] < num:
        return prefix_map[prefix]
    return num
nzkunb0c

nzkunb0c3#

我添加另一个答案,因为我的理解完全不同,所以方法也不同。您可以使用merge_asof

# Sort your dataframe and create the postal code prefix
df = df.sort_values('postalCode').assign(prefix=lambda x: x['postalCode'].astype(str).str.zfill(5).str[:2])

# Convert your zipcode list as dataframe and create the postal code prefix
zp = pd.DataFrame({'mainPostCode': zipcode}).sort_values('mainPostCode').assign(prefix=lambda x: x['mainPostCode'].astype(str).str.zfill(5).str[:2])

# Use merge_asof to merge on zip codes but first by prefix
# Look at the nearest lower value
out = (pd.merge_asof(df, zp, direction='backward', by='prefix',
                     left_on='postalCode', right_on='mainPostCode')
         .set_index(df.index).sort_index()
         .drop(columns='prefix').fillna(-1)
         .astype({'mainPostCode': int}))
print(out)

# Output
       postalCode  mainPostCode
0           23041         23039
1           48558         48556
2           52895         52893
3           39817         39808
4           40427         40427
...           ...           ...
39995       81184         81181
39996        7125          7122
39997       22773         22773
39998       88802         88799
39999       58510         58506
nwlqm0z1

nwlqm0z14#

预排序邮政总局邮政编码和使用二进制搜索。

from bisect import bisect_right

# MAKE SURE THEY ARE SORTED
zip_codes = list(range(1000, 100000))

# BEFORE
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 zip_codes if i > start and i < num], default=num)
    return m

# AFTER
def find_main_post_office(zip_code):
    start = int(str(zip_code // 100) + "00")
    index = bisect_right(zip_codes, start)
    try:
        return min(zip_codes[index], zip_code)
    except IndexError:
        return zip_code

"""
In [11]: %timeit findMainPostOffice(20050)
1.96 ms ± 1.33 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [16]: %timeit find_main_post_office(20050)
463 ns ± 0.462 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
"""

相关问题