add quadtree with quadtree:get_closest()

1 files changed, 50 insertions, 0 deletions

diff --git a/city/geometry.lua b/city/geometry.lua
index 410b4a3..b26779b 100644
--- a/city/geometry.lua
+++ b/city/geometry.lua
@@ -186,4 +186,54 @@ geom.qt_node = class{
    end
 }
 
+
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- quadtree
+--
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+geom.quadtree = class{
+   __call = function(this)
+      local tree = {}
+      tree.root = geom.qt_node(geom.point(0.5, 0.5), 1)
+
+      setmetatable(tree, {__index=this})
+      return tree
+   end,
+
+   insert = function(self, point)
+      self.root:insert(point)
+   end,
+
+   get_closest = function(self, point, accept_self)
+      local accept_self = accept_self or true
+
+      local region = self.root:find_containing(point)
+      local neighborhood = geom.square(point, region.square.span*4)
+
+      local near_points = {}
+      self.root:query(neighborhood, near_points)
+
+      if #near_points == 0 then return nil end
+      table.sort(near_points,
+		 function(a, b)
+		    return point:distance_to(a) < point:distance_to(b)
+		 end
+      )
+
+      local result
+
+      if not accept_self and
+	 (point.x == near_points[1].x) and
+	 (point.y == near_points[1].y) then
+	 result = near_points[2]
+      else
+	 result = near_points[1]
+      end
+
+      return result
+   end
+}
+
 return geom