Python日常之——不要再写循环了！

今天在做项目的时候遇到一个问题，给定一组点，如何方便快捷的求出点之间的距离。比如总共有100个点，两两组合求距离，总共有 $100\times99\div2=495$ 组距离。
如何求出这些距离呢？最先想到的是循环处理。
我们先定义距离函数：

import math
def dis_func(p1,p2):
    d = math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
    return d

循环处理，为了避免重复以及与点本身计算距离内循环从i + 1开始。：

distances = []
for i in range(len(points)):
    distance = []
    for j in range(i+1,len(points)):
        d = dis_func(points[i],points[j])
        distance.append(d)
    distances.append(distance)

结果是一个嵌套的阶梯状的列表，不容易使用，也不容易查找距离对应的点组是哪两个。于是可以这样写：

distances = []
for i in range(len(points)):
    distance = []
    for j in range(len(points)):
        if j <= i:
            d = None
        else:
            d = dis_func(points[i],points[j])
        distance.append(d)
    distances.append(distance)

我们填入了None作为占位符号。
最后用numpy库对嵌套列表转换：

distances = np.array(distances)

这样我们就很容易取出计算出的两个点的距离了，比如第23个点与第58个点的距离：

print(distances[22,57])

但如果点很多时由于python的运算速度所限，就不那么好用了。

3000个点距离计算时间

从结果来看，当数据点有3000个点时候，运算速度就已经很慢了，平均达到了4.73秒。
解决这个问题的办法是求助于第三方库，比如numpy:
假设目前就有3000个点，我们首先把点列表转换为numpy的array,然后分别取出来x与y值。

points = np.array(points)
points_xvalue = points[:,0]
points_yvalue = points[:,1]

再分别求 $dis = \sqrt{(x_1 - x_2)^2+(y_1 - y_2)^2}$ 中的 $(x_1 - x_2)^2$ 与 $(y_1 - y_2)^2$ 。

x1,x2 = np.meshgrid(points_xvalue,points_xvalue)
square_x = (x1 - x2)**2

y1,y2 = np.meshgrid(points_yvalue,points_yvalue)
square_y = (y1 - y2)**2

最后求结果：

res = np.sqrt(square_x + square_y)

可以将以上过程整合到一个函数里面：

def dis_func_new(xy):
    points_xvalue = points[:,0]
    points_yvalue = points[:,1]
    x1,x2 = np.meshgrid(points_xvalue,points_xvalue)
    square_x = (x1 - x2)**2
    y1,y2 = np.meshgrid(points_yvalue,points_yvalue)
    square_y = (y1 - y2)**2
    res = np.sqrt(square_x + square_y)
    return res

numpy数组3000个点距离计算

时间缩短到了0.4秒。但实际上鉴于numpy的广播机制，我们的代码还可以减少，运算速度也还可以再提升。

def dis_func_new(xy):
    x1,x2 = points[:,0],points[:,[0]]
    y1,y2 = points[:,1],points[:,[1]]
    square_x = (x1 - x2)**2
    square_y = (y1 - y2)**2
    res = np.sqrt(square_x + square_y)
    return res

这个函数代码少去了两行，运行结果也显示，时间也只0.26秒，性能进一步提升。
可见，在大数组运算时，numpy库相比python自有的嵌套列表循环，快了还是很多的！