python - 如何提高for循环的执行时间？

Question

我有一个名为“zone”的 DF，其中 x 和 y 列为整数，可以解释为点的位置。我需要计算第一个和第二个邻居的数量，我已经写了:

import numpy as np
import pandas as pd

data = np.random.randint(1000,6000,size=(600000,2))
zone = pd.DataFrame(data, columns=['x', 'y']).drop_duplicates()

a=[]
for i,row in zone.iterrows(): 
    x = row.x
    y = row.y
    num_1st_neigh = len(zone[(zone.x>=(x-1))&(zone.x<=(x+1))&(zone.y>=(y-1))&(zone.y<=(y+1))])-1
    num_2nd_neigh = (len(zone[(zone.x>=(x-2))&(zone.x<=(x+2))&(zone.y>=(y-2))&(zone.y<=(y+2))])-1)\
    -(num_1st_neigh)
    a.append([i,num_1st_neigh,num_2nd_neigh])
a = pd.DataFrame(a, columns = ['index','num_1st_neigh','num_2nd_neigh'])
zzz = zone.reset_index().merge(a,on='index')

这个效果很好，但在 3K 点上持续 15 秒，我有 1M 点，2 小时后它仍然运行。关于如何提高执行速度有什么想法吗？

我读到 iterrows 非常慢，但我不知道还能怎么做。

编辑:我也尝试使用 SQL 进行相同的操作，但执行时间>2小时并且查询返回超时:

SELECT t0.x,
    t0.y,
    count_if(greatest(abs(t0.x-t1.x), abs(t0.y-t1.y)) = 1) num_1_neighbors,
    count_if(greatest(abs(t0.x-t1.x), abs(t0.y-t1.y)) = 2) num_2_neighbors
FROM "table" t0 
    left join "table" t1 on t1.x between t0.x -2 and t0.x + 2
    and t1.y between t0.y -2 and t0.y + 2
    and (
        t1.x <> t0.x
        or t1.y <> t0.y
    )
group by 1,2

任何使用 SQL 或 pandas 的想法都非常受欢迎

Answer 1

您可以使用BallTree来自sklearn:

from sklearn.neighbors import BallTree

xy = zone[['x', 'y']]
tree = BallTree(xy, metric='euclidean')
num_1st_neigh = tree.query_radius(xy, r=1*np.sqrt(2), count_only=True) - 1
num_2nd_neigh = tree.query_radius(xy, r=2*np.sqrt(2), count_only=True) - 1 - num_1st_neigh
zone['num_1st_neigh'] = num_1st_neigh
zone['num_2nd_neigh'] = num_2nd_neigh

来自一个小例子:

# BallTree
>>> zone
     x    y  num_1st_neigh  num_2nd_neigh
0  106  115              0              0
1  118  104              0              0
2  119  114              0              0
3  108  103              0              2
4  103  101              0              0
5  110  105              0              1
6  103  104              0              0
7  102  119              0              0
8  106  105              0              1
9  111  114              0              0

# Your code
>>> zzz
   index    x    y  num_1st_neigh  num_2nd_neigh
0      0  106  115              0              0
1      1  118  104              0              0
2      2  119  114              0              0
3      3  108  103              0              2
4      4  103  101              0              0
5      5  110  105              0              1
6      6  103  104              0              0
7      7  102  119              0              0
8      8  106  105              0              1
9      9  111  114              0              0