quadtree.pyx

quadtree.pyx

"""
adapted from:
https://github.com/karimbahgat/Pyqtree/blob/master/pyqtree.py

TODO implement remove methods
"""

from libcpp cimport bool
from libcpp.vector cimport vector

ctypedef struct Rect:
    double x1
    double y1
    double x2
    double y2

ctypedef struct QuadNode:
    unsigned int id
    Rect rect

cdef struct QuadTree:
    # used so we can identify
    # uninitialized child quadtrees
    bool empty

    # center x, center y
    double x, y
    double width, height
    unsigned int depth
    vector[QuadNode] nodes
    QuadTree* children[4]

cdef QuadTree EMPTY = QuadTree(empty=True)


cdef Rect normalize_rect(double x1, double x2, double y1, double y2):
    if x1 > x2:
        x1, x2 = x2, x1
    if y1 > y2:
        y1, y2 = y2, y1
    return Rect(x1, y1, x2, y2)


cdef class QTree:
    cdef:
        QuadTree root
        vector[QuadTree] trees
        unsigned int max_items
        unsigned int max_depth

    def __init__(self, tuple bbox, unsigned int max_items=10, unsigned int max_depth=20):
        print('INITING')
        x1, y1, x2, y2 = bbox
        width, height = abs(x2-x1), abs(y2-y1)
        midx, midy = x1+width/2.0, y1+height/2.0

        self.root.empty = False
        self.root.children[0] = &EMPTY
        self.root.children[1] = &EMPTY
        self.root.children[2] = &EMPTY
        self.root.children[3] = &EMPTY
        self.root.x = midx
        self.root.y = midy
        self.root.width = width
        self.root.height = height
        self.max_items = max_items
        self.max_depth = max_depth

    cpdef intersect(self, bbox):
        cdef:
            set results = set()
            Rect rect = normalize_rect(bbox[0], bbox[1], bbox[2], bbox[3])
        return _intersect(&self.root, rect, results)

    cpdef insert(self, unsigned int id, tuple bbox):
        cdef Rect rect = normalize_rect(bbox[0], bbox[1], bbox[2], bbox[3])
        self._insert(&self.root, id, rect)

    cdef _insert(self, QuadTree *qtree, unsigned int id, Rect rect, int cdepth=0):
        cdef:
            QuadNode node
            int n

        rect = normalize_rect(rect.x1, rect.y1, rect.x2, rect.y2)

        # TODO is this the correct
        # way to check if the array is "empty"?
        if qtree.children[0].empty:
            node = QuadNode(id, rect)
            n = qtree.nodes.size()
            print('{}NODES SIZE {}'.format('  ' * cdepth, n))
            qtree.nodes.push_back(node)

            if qtree.nodes.size() > self.max_items and qtree.depth < self.max_depth:
                self._split(qtree, cdepth+1)
        else:
            self._insert_into_children(qtree, id, rect, cdepth+1)

    cdef _insert_into_children(self, QuadTree *qtree, unsigned int id, Rect rect, int cdepth):
        cdef QuadNode node

        # if rect spans center then insert here
        if (rect.x1 <= qtree.x and rect.x2 >= qtree.x and
            rect.y1 <= qtree.y and rect.y2 >= qtree.y):
            node = QuadNode(id, rect)
            qtree.nodes.push_back(node)
        else:
            # try to insert into children
            if rect.x1 <= qtree.x:
                if rect.y1 <= qtree.y:
                    self._insert(qtree.children[0], id, rect, cdepth+1)
                if rect.y2 >= qtree.y:
                    self._insert(qtree.children[1], id, rect, cdepth+1)
            if rect.x2 > qtree.x:
                if rect.y1 <= qtree.y:
                    self._insert(qtree.children[2], id, rect, cdepth+1)
                if rect.y2 >= qtree.y:
                    self._insert(qtree.children[3], id, rect, cdepth+1)

    cdef _split(self, QuadTree *qtree, int cdepth):
        cdef:
            double quartwidth = qtree.width / 4.0
            double quartheight = qtree.height / 4.0
            double halfwidth = qtree.width / 2.0
            double halfheight = qtree.height / 2.0
            double x1 = qtree.x - quartwidth
            double x2 = qtree.x + quartwidth
            double y1 = qtree.y - quartheight
            double y2 = qtree.y + quartheight
            unsigned int new_depth = qtree.depth + 1
            QuadTree q1 = QuadTree(x=x1, y=y1, width=halfwidth, height=halfheight, depth=new_depth, empty=False, nodes=[])
            QuadTree q2 = QuadTree(x=x1, y=y2, width=halfwidth, height=halfheight, depth=new_depth, empty=False, nodes=[])
            QuadTree q3 = QuadTree(x=x2, y=y1, width=halfwidth, height=halfheight, depth=new_depth, empty=False, nodes=[])
            QuadTree q4 = QuadTree(x=x2, y=y2, width=halfwidth, height=halfheight, depth=new_depth, empty=False, nodes=[])
            vector[QuadNode] nodes

        self.trees.push_back(q1)
        self.trees.push_back(q2)
        self.trees.push_back(q3)
        self.trees.push_back(q4)

        q1.children[0] = &EMPTY
        q1.children[1] = &EMPTY
        q1.children[2] = &EMPTY
        q1.children[3] = &EMPTY

        q2.children[0] = &EMPTY
        q2.children[1] = &EMPTY
        q2.children[2] = &EMPTY
        q2.children[3] = &EMPTY

        q3.children[0] = &EMPTY
        q3.children[1] = &EMPTY
        q3.children[2] = &EMPTY
        q3.children[3] = &EMPTY

        q4.children[0] = &EMPTY
        q4.children[1] = &EMPTY
        q4.children[2] = &EMPTY
        q4.children[3] = &EMPTY

        # ...will this cause a memory leak?
        qtree.children = [&q1, &q2, &q3, &q4]

        nodes.swap(qtree.nodes)
        for node in nodes:
            self._insert_into_children(qtree, node.id, node.rect, cdepth+1)



cdef _intersect(QuadTree *qtree, Rect rect, set results):
    # search children
    # TODO is this the correct
    # way to check if the array is "empty"?
    if not qtree.children[0].empty:
        if rect.x1 <= qtree.x:
            if rect.y1 <= qtree.y:
                _intersect(qtree.children[0], rect, results)
            if rect.y2 >= qtree.y:
                _intersect(qtree.children[1], rect, results)
        if rect.x2 >= qtree.x:
            if rect.y1 <= qtree.y:
                _intersect(qtree.children[2], rect, results)
            if rect.y2 >= qtree.y:
                _intersect(qtree.children[3], rect, results)

    # search node at this level
    for node in qtree.nodes:
        if (node.rect.x2 >= rect.x1 and node.rect.x1 <= rect.x2 and
            node.rect.y2 >= rect.y1 and node.rect.y1 <= rect.y2):
            results.add(node.id)
    return results