kd-tree Visualization(3)

さて、先日のエントリにて、メディアンを加味した2d木の構築を紹介しました。コードの見通しを簡潔にする目的もあり、メディアンの選択にはsort()を用いました。

#!/usr/bin/python -t

from operator import itemgetter
from random import seed, randint

def build(points, d=1):
    pts = list(points)

    if not pts:
        return None
    
    h = len(pts) // 2

    pts.sort(key=itemgetter(d))

    d = (d + 1) % 2

    return [build(pts[0:h],  d),
            build(pts[h+1:], d),
            pts[h]]

if __name__ == '__main__':
    seed(0)
    pts = [(randint(1, 100), randint(1, 100))
           for i in range(32)]

    tree = [None, build(pts), (0, 0)]

一般的にsort()を用いたメディアンの選択はコストが高いことで知られています。また、kd木は再帰的に木構造を構築するアルゴリズムです。そのため、均衡の取れた木構造を構築するためには、メディアンも再帰的に分割される領域の数とほぼ同じ回数分選択しなければなりません。これは、サンプルの数が増えれば増えるほど、sort()に費やすコストが高く付くことを意味します。

さて、少々メディアンから離れて、座標群の整列について考えてみましょう。座標群は予め整列していなければならないでしょうか? その必要はありません*1。検証も兼ねて、先日のエントリにて記載したコードを変更してみましょう。

from operator import itemgetter
from itertools import groupby

def median(iterable, key=None):
    if key is None:
        key = lambda x: x
    pool = tuple(iterable)
    h = len(pool) // 2
    return sorted(pool, key=key)[h]

def build(points, d=1):
    pts = tuple(points)

    if not pts:
        return None

    m = median(pts, key=itemgetter(d))

    dpts = { -1: [], 0: [], 1: [], }
    for k, g in groupby(pts, key=lambda x: (x[d] > m[d]) - (x[d] < m[d])):
        dpts[k].extend(g)

    h = len(dpts[0]) // 2

    d = (d + 1) % 2

    return [build(dpts[-1] + dpts[0][0:h],    d),
            build(dpts[ 1] + dpts[0][h+1:-1], d),
            dpts[0][h]]