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函数中计算30k * 30k = 9亿个距离，即纯Python。

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将这些带入连接中的longint;我不知道如何解决这个问题），这可以达到大约20个字节* 900m行= 18GB用于最终的数据框架。它在实践中可能更多，并且在操作期间的峰值内存使用率高于此值（特别是，排序将花费最长时间，并使用大量内存）。

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

在一天结束时，你正在研究接近十亿计算的计算;如果你想以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()方法进行一些堆叠/取消堆叠，您可以大大提高速度，但在这种情况下，由于您的数据包含重复项，因此它有点棘手。）

在熊猫中更快的应用方法

1 个答案: