3.10.1 计算最近旁点

众所周知，对于一维的已排序的数值，可以使用二分法快速找到与指定数值最近的数值。在下面的例子中，在一个升序排序的随机数数组x中，使用numpy.searchsorted()搜索离0.5最近的数。排序算法的时间复杂度为O(NlogN)，而每次二分搜索的时间复杂度为O(logN)。

    x = np.sort(np.random.rand(100))
    idx = np.searchsorted(x, 0.5)
    print x[idx], x[idx - 1] #距离0.5最近的数是这两个数中的一个
    0.542258714465 0.492205345391

类似的算法可以推广到N维空间，spatial模块提供的cKDTree类使用K-d树快速搜索N维空间中的最近点。在下面的例子中，用平面上随机的100个点创建cKDTree对象，并对其搜索与targets中每个点距离最近的3个点。cKDTree.query()返回两个数组dist和idx，dist[i, :]是距离targets[i]最近的3个点的距离，而idx[i, :]是这些最近点的下标：

    from scipy import spatial
    np.random.seed(42)
    N = 100
    points = np.random.uniform(-1, 1, (N, 2))
    kd = spatial.cKDTree(points)
    
    targets = np.array([(0, 0), (0.5, 0.5), (-0.5, 0.5), (0.5, -0.5), (-0.5, -0.5)])
    dist, idx = kd.query(targets, 3)
                       dist                         idx
    -----------------------------------------  --------------
    [[ 0.15188266,  0.21919416,  0.27647793],  [[48, 73, 81],
     [ 0.09595807,  0.15745334,  0.22855398],   [37, 78, 43],
     [ 0.05009422,  0.17583445,  0.1807312 ],   [79, 22, 92],
     [ 0.11180181,  0.16618122,  0.18127473],   [35, 58,  6],
     [ 0.19015485,  0.19060739,  0.19361173]]   [83,  7, 42]]

cKDTree.query_ball_point()搜索与指定点在一定距离之内的所有点，它只返回最近点的下标，由于每个目标点的近旁点数不一定相同，因此idx2数组中的每个元素都是一个列表：

    r = 0.2
    idx2 = kd.query_ball_point(targets, r)
    idx2
    array([[48], [37, 78], [79, 92, 22], [58, 35, 6], [7, 55, 83, 42]], dtype=object)

cKDTree.query_pairs()找到points中距离小于指定值的每一对点，它返回的是一个下标对的集合对象。下面的程序使用集合差集运算找出所有距离在0.08到0.1之间的点对：

    idx3 = kd.query_pairs(0.1) - kd.query_pairs(0.08)
    idx3
    {(1, 46),     (3, 21),     (3, 82),     (3, 95),     (5, 16),     (9, 30),
     (10, 87),    (11, 42),    (11, 97),    (18, 41),    (29, 74),    (32, 51),
     (37, 78),    (39, 61),    (41, 61),    (50, 84),    (55, 83),    (73, 81)}

在图3-57中，与target中的每个点(用五角星表示)最近的点用与其相同的颜色标识：

图3-57 用cKDTree寻找近旁点

cKDTree的所有搜索方法都有一个参数p，用于定义计算两点之间距离的函数。读者可以尝试使用不同的p参数，观察图3-57的变化：

●p=1：绝对值之和作为距离

●p=2：欧式距离

●p=np.inf：最大坐标差值作为距离

此外，cKDTree.query_ball_tree()可以在两棵K-d树之间搜索距离小于给定值的所有点对。

distance子模块中的pdist()计算一组点中每对点的距离，而cdist()计算两组点中每对点的距离。由于pdist()返回的是一个压缩之后的一维数组，需要用squareform()将其转换成二维数组。dist1[i, j]是points中下标为i和j的两个点的距离，dist2[i, j]是points[i]和targets[j]之间的距离。

    from scipy.spatial import distance
    dist1 = distance.squareform(distance.pdist(points))
    dist2 = distance.cdist(points, targets)
    dist1.shape  dist2.shape
    -----------  -----------
    (100, 100)   (100, 5)

下面使用np.min()在dist2中搜索points中与targets距离最近的点，其结果与cKDTree.query()的结果相同：

    print dist[:, 0] # cKDTree.query()返回的与targets最近的距离
    print np.min(dist2, axis=0)
    [ 0.15188266  0.09595807  0.05009422  0.11180181  0.19015485]
    [ 0.15188266  0.09595807  0.05009422  0.11180181  0.19015485]

为了找到points中最近的点对，需要将dist1对角线上的元素填充为无穷大：

    dist1[np.diag_indices(len(points))] = np.inf
    nearest_pair = np.unravel_index(np.argmin(dist1), dist1.shape)
    print nearest_pair, dist1[nearest_pair]
    (22, 92) 0.00534621024816

用cKDTree.query()可以快速找到这个距离最近的点对：在K-d树中搜索它自己包含的点，找到与每个点最近的两个点，其中距离最近的点就是它本身，距离为0，而距离第二近的点就是每个点的最近旁点，然后只需要找到这些距离中最小的那个即可：

    dist, idx = kd.query(points, 2)
    print idx[np.argmin(dist[:, 1])], np.min(dist[:, 1])
    [22 92] 0.00534621024816

让我们看一个使用K-d树提高搜索速度的实例。下面的start和end数组保存用户登录和离开网站的时间，对于任意指定的时刻time，计算该时刻在线用户的数量。

    N = 1000000
    start = np.random.uniform(0, 100, N)
    span = np.random.uniform(0.01, 1, N)
    span = np.clip(span, 2, 100)
    end = start + span

下面的naive_count_at()采用逐个比较的方法计算指定时间的在线用户数量：

    def naive_count_at(start, end, time):
        mask = (start＜time)＆(end＞time)
        return np.sum(mask)

图3-58显示了如何使用K-d树实现快速搜索。图中每点的横坐标为start，纵坐标为end。由于end大于start，因此所有的点都在y=x斜线的上方。图中阴影部分表示满足(start＜time)＆(end＞time)条件的区域。该区域中的点数为时刻time时的在线人数。

图3-58 使用二维K-d树搜索指定区间的在线用户

tree.count_neighbors(other, r, p=2)可以统计tree中到other的距离小于r的点数，其中p为计算距离的范数。距离按照如下公式计算：

当p=2时，上面的公式就是欧几里得空间中的向量长度。当p=∞时，该距离变为各个轴上的最大绝对值：

N∞ (x)=｜｜x｜｜∞ =max(｜x1｜,｜x2｜…,｜xn｜)

当使用p=∞时可以计算某正方形区域之内的点数。我们将该正方形的中心设置为(time - max_time, time + max_time)，正方形的边长为2 * max_time，即r=max_time。其中max_time为end中的最大值。下面的KdSearch类实现了该算法：

    class KdSearch(object):
        def __init__(self, start, end, leafsize=10):
            self.tree = spatial.cKDTree(np.c_[start, end], leafsize=leafsize)
            self.max_time = np.max(end)
    
        def count_at(self, time):
            max_time = self.max_time
            to_search = spatial.cKDTree([[time - max_time, time + max_time]])
            return self.tree.count_neighbors(to_search, max_time, p=np.inf)
    
    naive_count_at(start, end, 40) == KdSearch(start, end).count_at(40)
    True

下面比较运算时间，由结果可知创建K-d树需要约0.5秒时间，K-d树的搜索速度则为线性搜索的17倍左右。

请读者研究点数N和leafsize参数与创建K-d树和搜索时间之间的关系。

    %time ks = KdSearch(start, end, leafsize=100)
    %timeit naive_search(start, end, 40)
    %timeit ks.count_at(40)
    Wall time: 484 ms
    100 loops, best of 3: 3.85 ms per loop
    1000 loops, best of 3: 221 µs per loop
上一章目录下一章