Quad tree implement using bin tree. Uses pygame to display.

quadtree.py

#!/usr/bin/env python2
import pygame
from pygame.locals import *

COLOR_BG = (0xff, 0xff, 0xff, 0xff)
COLOR_BLACK = (0, 0, 0, 0xff)
COLOR_TRANS = (0, 0, 0, 0)
G = 1

class SMTNode:
    def __init__(self, val):
        self.val = val
        self.lch = self.rch = None

class Rect:
    HSPLIT, VSPLIT = 0, 1
    def __init__(self, x1, y1, x2, y2):
        if not (x1 <= x2 and y1 <= y2):
            raise Exception("invalid rect (%d, %d, %d, %d)"%(x1, y1, x2, y2))
        self.x1 = x1
        self.x2 = x2
        self.y1 = y1
        self.y2 = y2
    def __eq__(self, r):
        return (self.x1 == r.x1 and self.x2 == r.x2 and self.y1 == r.y1 and self.y2 == r.y2)
    def __repr__(self):
        return '(%r, %r)'%((self.x1, self.y1), (self.x2, self.y2))
    def contains(self, r):
        return (self.x1 <= r.x1 and self.x2 >= r.x2 and self.y1 <= r.y1 and self.y2 >= r.y2)
    def intersection(self, r):
        return Rect(max(self.x1, r.x1), max(self.y1, r.y1), min(self.x2, r.x2), min(self.y2, r.y2))
    def split(self):
        x1, y1, x2, y2 = self.x1, self.y1, self.x2, self.y2
        dx = x2 - x1
        dy = y2 - y1
        if dx > dy:
            mid = (x1 + x2) // 2
            return Rect.VSPLIT, Rect(x1, y1, mid, y2), Rect(mid, y1, x2, y2)
        else:
            mid = (y1 + y2) // 2
            return Rect.HSPLIT, Rect(x1, y1, x2, mid), Rect(x1, mid, x2, y2)

class SMTree:
    def __init__(self, r, val, surface):
        self.rect = r
        self.root = SMTNode(val)
        self.root.cnt = 1
        self.surface = surface
        self.draw_rect(r, (0xee, 0xee, 0xee))

    def _insert(self, p, tr, r, val):
        # print 'insert %s to %s with %s' %(r, tr, val)
        if tr == r:
            p.val = val
            p.lch = p.rch = None
            self.draw_rect(tr, val)
            return
        dir, r1, r2 = tr.split()
        if p.lch is None:
            p.lch = SMTNode(p.val)
            p.lch.cnt = 1
            p.rch = SMTNode(p.val)
            p.rch.cnt = 1
            p.val = None
        # draw the splited line
        if dir == Rect.HSPLIT:
            self.draw_line((r1.x1, r1.y2-1), (r1.x2, r1.y2-1), COLOR_BLACK)
            self.draw_line((r2.x1, r1.y1), (r2.x2, r1.y1), COLOR_BLACK)
        else:
            self.draw_line((r1.x2-1, r1.y1), (r1.x2-1, r1.y2-1), COLOR_BLACK)
            self.draw_line((r2.x1, r2.y1), (r2.x1, r2.y2-1), COLOR_BLACK)
        if r1.contains(r):
            self._insert(p.lch, r1, r, val)
        elif r2.contains(r):
            self._insert(p.rch, r2, r, val)
        else:
            self._insert(p.lch, r1, r1.intersection(r), val)
            self._insert(p.rch, r2, r2.intersection(r), val)
        p.cnt = 1 + p.lch.cnt + p.rch.cnt

    def insert(self, r, val):
        self._insert(self.root, self.rect, self.rect.intersection(r), val)
        print "nodes %d" % self.root.cnt

    def query(self, r):
        pass

    def draw_rect(self, r, color):
        if isinstance(r, Rect):
            pygame.draw.rect(self.surface, color, (r.x1, r.y1, r.x2-r.x1, r.y2-r.y1))
        else:
            raise Exception("r: excepted: Rect, got: %s" %(type(r)))

    def draw_line(self, p1, p2, color, width=1):
        p1 = int(p1[0]), int(p1[1])
        p2 = int(p2[0]), int(p2[1])
        # print 'draw line %s-%s' % (p1, p2)
        pygame.draw.line(self.surface, color, p1, p2, width) 

class Drawer:
    def __init__(self, surface):
        self.state = 0
        self.surface = surface
    def update(self, e):
        if self.state == 0:
            if e.type == MOUSEBUTTONDOWN:
                self.state = 1
                self.start = e.pos
        elif self.state == 1:
            if e.type == MOUSEBUTTONUP:
                self.state = 2
                self.end = e.pos
            elif e.type == MOUSEMOTION:
                self.end = e.pos
                x1, y1 = self._convert(self.start)
                x2, y2 = self._convert(self.end)
                self.surface.fill(COLOR_TRANS)
                pygame.draw.rect(self.surface, (0, 0x55, 0, 0xaa), (x1, y1, x2-x1, y2-y1))
    def _convert(self, p):
        x, y = p
        return x//G*G, y//G*G

    def get_result(self):
        if self.state != 2: 
            return None
        x1, y1 = self._convert(self.start)
        x2, y2 = self._convert(self.end)
        r = Rect(min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2))
        self.state = 0
        self.surface.fill(COLOR_TRANS)
        return r

if __name__ == '__main__':
    # init pygame stuff
    W, H = 900, 700
    pygame.display.init()
    screen = pygame.display.set_mode((W, H), 0, 32)
    # configure colors
    c0, c1 = (0xff, 0xff, 0xff, 0xff), (0x95, 0, 0, 0xff)
    # create graphical layers
    layer1 = pygame.Surface((W, H)).convert_alpha()
    layer1.fill(COLOR_TRANS)
    layer2 = pygame.Surface((W, H)).convert_alpha()
    layer2.fill(COLOR_TRANS)
    # create the quad tree, the tree draws on layer1
    # tree = SMTree(Rect(0, 0, 35 * G, 34 * G), c0, layer1)
    tree = SMTree(Rect(0, 0, 813 * G, 613 * G), c0, layer1)
    # drawer use layer2 to draw
    drawer = Drawer(layer2)
    # start main loop
    tm = pygame.time.Clock()
    while 1:
        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                exit(0)
            drawer.update(e)
        r = drawer.get_result()
        if r:
            tree.insert(r, c1)
        screen.fill(COLOR_BG)
        screen.blit(layer1, (0, 0))
        screen.blit(layer2, (0, 0))

        pygame.display.flip()
        tm.tick(30)