PythonでQuadtree - 2


[前回](http://qiita.com/yamaguchiyuto/items/de60d8fe400deeabbd97)の実装だとある点がどの領域に属すかを調べるのにO(領域数）かかってしまう事に気づいたので実装しなおした

今回はちゃんと木構造を保持しつつ分割するようにした

こうすればあるデータ点がどの領域に属すかを調べるのをO（木の深さ）でできる

というかこれが普通

##実装

同一のデータ点が複数あると無限に分割を繰り返してしまう可能性があるので、最大の分割数（`maxdivision`）を与えている

```quadtree.py
# -*- coding: utf-8 -*-
class Quadtree:
    def __init__(self,data,x1,y1,x2,y2,maxpoints,maxdivision):
        self.data = data
        self.sequential_id = 0
        self.leaves = {}
        self.root = self.divide(self.init_area(data,x1,y1,x2,y2),maxpoints,maxdivision)

    def __iter__(self):
        for aid in self.leaves:
            yield self.leaves[aid]

    def init_area(self,data,x1,y1,x2,y2):
        initial = Area(x1,y1,x2,y2)
        for d in data:
            initial.append(d)
        return initial

    def subdivide(self,area):
        division = []

        """ 分割後の各辺の長さ """
        xl = (area.x2 - area.x1)/2
        yl = (area.y2 - area.y1)/2

        """ 分割後の領域を生成 """
        for dx in [0,1]:
            for dy in [0,1]:
                """
                0 2
                1 3
                の順番
                """
                sub_area = Area(area.x1+dx*xl, area.y1+dy*yl, area.x1+(1+dx)*xl, area.y1+(1+dy)*yl)
                division.append(sub_area)

        """ 分割前の領域に属すデータ点を分割後の領域にアサイン """
        for p in area.points():
            for sub_area in division:
                if sub_area.cover(p):
                    sub_area.append(p)
                    break

        return division

    def divide(self, area, maxpoints, division_left):
        """ 引数で与えられたdivision_left回だけ分割を繰り返す """
        if division_left == 0 or area.number_of_points() <= maxpoints:
            """ areaに含まれるデータ点の数がmaxpointsを超えていなければ分割終了 """
            area.set_fixed(self.sequential_id)
            area.calc_center() # areaに属すデータ点のセントロイドを計算
            self.leaves[self.sequential_id] = area
            self.sequential_id += 1
            return area
        else:
            """ areaに含まれるデータ点の数がmaxpointsを超えていれば四分割 """
            next_level = self.subdivide(area)
            """ 子エリアそれぞれを更に分割 """
            for i in range(4):
                child = self.divide(next_level[i],maxpoints,division_left-1)
                area.set_child(i,child)
            """ 分割後のエリアを返す """
            return area

    def get_area_id(self,p):
        return self.root.covered(p)

class Area:
    def __init__(self,x1,y1,x2,y2):
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2
        self.points_ = []
        self.fixed = False
        """
        0 2
        1 3
        """
        self.children = [None,None,None,None]

    def calc_center(self):
        if self.number_of_points() == 0:
            self.center = (self.x1 + (self.x2-self.x1)/2, self.y1 + (self.y2-self.y1)/2)
        else:
            sumx = 0.
            sumy = 0.
            for p in self.points_:
                sumx += p[0]
                sumy += p[1]
            self.center = (sumx/len(self.points_), sumy/len(self.points_))

    def append(self, p):
        """ 領域にデータ点を追加 """
        self.points_.append(p)

    def points(self):
        """ 領域に属しているデータ点を返す """
        return self.points_

    def number_of_points(self):
        return len(self.points_)

    def is_fixed(self):
        """ 分割が終わっているかどうか """
        return self.fixed

    def set_fixed(self,aid):
        """ 分割が終わったフラグを立てる """
        self.fixed = True
        """ エリアにIDを付ける """
        self.aid = aid

    def set_child(self,n,area):
        self.children[n] = area

    def covered(self,p):
        if self.cover(p):
            if self.fixed:
                return self.aid
            else:
                """ 子エリアのIDを計算 """
                cid = 0
                if self.x1 + (self.x2 - self.x1) / 2 < p[0]:
                    """ Right """
                    cid += 2
                if self.y1 + (self.y2 - self.y1) / 2 < p[1]:
                    """ Down """
                    cid += 1
                return self.children[cid].covered(p)
        else:
            return None

    def cover(self, p):
        """ あるデータ点pがこの領域にカバーされるかどうか """
        if self.x1 < p[0] and self.y1 < p[1] and self.x2 >= p[0] and self.y2 >= p[1]:
            return True
        else:
            return False

if __name__ == '__main__':
    import sys

    def read_data(file_name):
        data = []
        for line in open(file_name, 'r'):
            entries = line.rstrip().split(' ')
            lat = float(entries[0])
            lng = float(entries[1])
            data.append((lat,lng))
        return data

    x1 = float(sys.argv[1])
    y1 = float(sys.argv[2])
    x2 = float(sys.argv[3])
    y2 = float(sys.argv[4])
    maxpoints = int(sys.argv[5])
    maxdivision = int(sys.argv[6])
    data = read_data(sys.argv[7])

    """ 分割 """
    qtree = Quadtree(data,x1,y1,x2,y2,maxpoints,maxdivision)

    """ 結果 """
    for a in qtree:
        print a.aid,a.x1,a.y1,a.x2,a.y2,a.center

    p = (0.37,0.55)
    print p,qtree.get_area_id(p)
```

##実行例

引数は対象とする領域の左上の座標のx、y、右下の座標のx、y、領域に属すことを許す最大のデータ数、最大分割数、データファイルの順で与える

```
$ cat data
0.567603099626 0.410160220857
0.405568042222 0.555583016695
0.450289054963 0.252870772505
0.373657009068 0.549501477427
0.500192599714 0.352420542886
0.626796922 0.422685113179
0.527521290061 0.483502904656
0.407737520746 0.570049935936
0.578095278433 0.6959689898
0.271957975882 0.450460115198
0.56451369371 0.491139311353
0.532304354436 0.486931064003
0.553942716039 0.51953331907
0.627341495722 0.396617894317
0.454210189397 0.570214499065
0.327359895038 0.582972137899
0.422271372537 0.560892624101
0.443036148978 0.434529240506
0.644625936719 0.542486338813
0.447813648487 0.575896033203
0.534217713171 0.636123087401
0.348346109137 0.312959224746
0.484075186327 0.289521849258
0.588417643962 0.387831556678
0.613422176662 0.665770829308
0.60994411786 0.399778040078
0.425443751505 0.402619561402
0.504955932504 0.610015349003
0.402852203978 0.382379275531
0.387591801531 0.452180343665

$ python quadtree 0 0 1 1 3 3 data
0 0.0 0.0 0.25 0.25 (0.125, 0.125)
1 0.0 0.25 0.25 0.5 (0.125, 0.375)
2 0.25 0.0 0.5 0.25 (0.375, 0.125)
3 0.25 0.25 0.375 0.375 (0.348346109137, 0.312959224746)
4 0.25 0.375 0.375 0.5 (0.271957975882, 0.450460115198)
5 0.375 0.25 0.5 0.375 (0.467182120645, 0.2711963108815)
6 0.375 0.375 0.5 0.5 (0.414730976498, 0.417927105276)
7 0.0 0.5 0.25 0.75 (0.125, 0.625)
8 0.0 0.75 0.25 1.0 (0.125, 0.875)
9 0.25 0.5 0.375 0.625 (0.35050845205299996, 0.566236807663)
10 0.25 0.625 0.375 0.75 (0.3125, 0.6875)
11 0.375 0.5 0.5 0.625 (0.42752015467779997, 0.5665272218)
12 0.375 0.625 0.5 0.75 (0.4375, 0.6875)
13 0.25 0.75 0.5 1.0 (0.375, 0.875)
14 0.5 0.0 0.75 0.25 (0.625, 0.125)
15 0.5 0.25 0.625 0.375 (0.500192599714, 0.352420542886)
16 0.5 0.375 0.625 0.5 (0.5650506999425, 0.44322384960416666)
17 0.625 0.25 0.75 0.375 (0.6875, 0.3125)
18 0.625 0.375 0.75 0.5 (0.627069208861, 0.40965150374799997)
19 0.75 0.0 1.0 0.25 (0.875, 0.125)
20 0.75 0.25 1.0 0.5 (0.875, 0.375)
21 0.5 0.5 0.625 0.625 (0.5294493242714999, 0.5647743340365)
22 0.5 0.625 0.625 0.75 (0.5752450560886667, 0.6659543021696667)
23 0.625 0.5 0.75 0.625 (0.644625936719, 0.542486338813)
24 0.625 0.625 0.75 0.75 (0.6875, 0.6875)
25 0.5 0.75 0.75 1.0 (0.625, 0.875)
26 0.75 0.5 1.0 0.75 (0.875, 0.625)
27 0.75 0.75 1.0 1.0 (0.875, 0.875)
(0.37, 0.55) 9
```

結果として28個の領域（木のリーフノードの数）に分割された。

結果の一番最後の行を見ると、与えられたデータ点がどの領域に属すかを正しく求められていることが分かる。