Something like a QuadTree for space partitioning

quadtree.py

#!/usr/bin/env python

class TreeNode(object):
  def __init__(self, parent, x_min, y_min, x_max, y_max, items = []):
    (self.x_min, self.y_min) = (x_min, y_min)
    (self.x_max, self.y_max) = (x_max, y_max)

    self.x_mid = (x_min + x_max) / 2
    self.y_mid = (y_min + y_max) / 2

    self.parent = parent
    self.items = items
    self.nw = None
    self.ne = None
    self.sw = None
    self.se = None

  def is_leaf(self):
    if self.nw: return False
    else: return True

  def divide(self):
    self.nw = TreeNode(self, x_min, y_min, x_mid, y_mid)
    self.ne = TreeNode(self, x_mid, y_min, x_max, y_mid)
    self.sw = TreeNode(self, x_min, y_mid, x_mid, y_max)
    self.se = TreeNode(self, x_mid, y_mid, x_max, y_max)
    for p in items:
      if p.x < x_mid: # West
        if p.y < y_mid: # North
          self.nw.items.push(p)
        else: # South
          self.sw.items.push(p)
      else: # East
        if p.y < y_mid: # North
          self.ne.items.push(p)
        else: # South
          self.se.items.push(p)
    self.items = None

  def neighbors(self, p):
    quadrants = []
    if p.x < (self.x_mid + p.radius) and p.y < (self.y_mid + p.radius): # North West
      quadrants.push(self.nw)
    if p.x > (self.x_mid - p.radius) and p.y < (self.y_mid + p.radius): # North East
      quadrants.push(self.ne)
    if p.x < (self.x_mid + p.radius) and p.y > (self.y_mid - p.radius): # South West
      quadrants.push(self.sw)
    if p.x > (self.x_mid - p.radius) and p.y > (self.y_mid - p.radius): # South East
      quadrants.push(self.se)
    return quadrants


  def search(self, p):
    if self.is_leaf():
      return self.items
    else:
      res = []
      for q in self.neighbors(p):
        res += q.search(p)
      return res