熊猫中更快的应用方法

Question

我有一个函数正在尝试应用于位置数据框。具体来说，我想添加一个新列，其中包含与每个站点最近的 10 个站点。以下似乎有效，但速度非常慢。

def distance(first_lat, first_lon, second_lat, second_lon):
    return ((first_lat - second_lat) ** 2 + (first_lon - second_lon) ** 2) ** 0.5


def load_site_list():
    '''
    This function generates a dataframe with all the available sites
    '''
    url = 'ftp://ftp.ncdc.noaa.gov/pub/data/noaa/isd-history.csv'
    cols = ["STATION NAME",
            "LAT",
            "LON"]
    df = pd.read_csv(url, parse_dates=False, usecols=cols)
    df = df.dropna(subset=['LAT'])
    df = df.dropna(subset=['LON'])
    df['LAT'] = df['LAT'].astype(float)
    df['LON'] = df['LON'].astype(float)
    return df

sites = load_site_list()
sites['closest'] = ""
for index, row in sites.iterrows():
    sites['dist'] = sites.apply(lambda line: distance(line['LAT'], line['LON'], row['LAT'], row['LON']), axis=1)
    sites.sort_values('dist', inplace=True)
    sites['closest'][index] = sites['STATION NAME'].iloc[1:11].tolist()

似乎 for 循环中生成与当前列的距离的第一行在每个循环中占用了第二个。这里有超过 10,000 行要循环...有没有更快的方法？

Answer 1

请注意，您的代码的时间复杂度为 O(n^2)：在这种情况下，您在 for 循环中的 apply 函数（即纯 Python）内计算 30k*30k=9 亿距离。

pandas 中的向量运算是在 C 中实现的，因此如果您在单个向量运算中计算所有距离，您将获得相对加速。

如果您有足够的 RAM，您可以进行笛卡尔连接，计算所有成对距离，然后进行排序、分组，然后取头，如下所示：

# code to reduce memory usage
sites['site_code'] = pd.Categorical(sites['STATION NAME']).codes
sites['LAT'] = sites.LAT.astype(np.float16)
sites['LON'] = sites.LAT.astype(np.float16)
sites_small = sites[['site_code','LAT','LON']].copy()
sites_small.index = [0]*len(sites_small)

pairs = sites_small.join(sites_small,lsuffix='_x',rsuffix='_y')
pairs['dist'] = (pairs['LAT_x'] - pairs['LAT_y'])**2 + (pairs['LON_x'] - pairs['LON_y'])**2
pairs.sort_values(['STATION NAME_x','dist'], inplace = True) # actually, just sorting by dist is sufficient
pairs.groupby('STATION NAME_x').head(10)

不幸的是，您可能没有足够的 RAM：如果您将站点名称编码为 16 位整数，并将坐标编码为 16 位浮点数，则每行需要 12 个字节（因为您正在查看对） ，再加上另外 8 个字节的索引（pandas 将这些放入连接中的 longints；我不知道如何解决这个问题），最终数据帧大约需要 20 个字节 * 900m 行 = 18GB。在实践中可能更多，并且操作期间的峰值内存使用率高于此（特别是排序将花费最长的时间，并且使用大量内存）。

我在我的机器上试过这个：我使用了大约 30GB，放弃了等待完整排序，而是对 dist 小于 100 的子集进行排序。用了不到 5 分钟，大部分时间都花在了加入上。

归根结底，您的计算量接近 10 亿次；如果您想以 C 的速度执行此操作而不必存储所有成对数据（在 pandas 中使用直接方法就是这种情况），您很可能必须使用 numpy 数组和/或使用 Cython 编写代码多处理。

更聪明的方法是避免进行十亿次计算，这涉及到知道哪些距离不需要计算。这需要一些巧妙的逻辑，但幸运的是，这是一个经过充分研究的 k-Nearest Neighbors 主题，它具有专门针对此类问题设计的高效算法：

from sklearn.neighbors import NearestNeighbors
data = sites[['LAT','LON']].values
nbrs = NearestNeighbors(n_neighbors=10, algorithm='auto', metric = 'euclidean').fit(data)
distances, indices = nbrs.kneighbors(data)
indices

这需要不到一秒钟的时间来计算。恢复最近邻居的名字需要更长的时间：

df = pd.DataFrame(indices, index = sites['STATION NAME'].values)
df.replace(dict(enumerate(sites['STATION NAME'].values)), inplace = True)

（实际上，您可以通过使用带有一些堆叠/取消堆叠的.merge() 方法来大大加快这一速度，但在这种情况下，它会稍微复杂一些，因为您的数据包含重复项。）