add quadtree with quadtree:get_closest()

2 files changed, 108 insertions, 0 deletions

diff --git a/city/geometry-test.lua b/city/geometry-test.lua
index 98275a3..28d55bc 100644
--- a/city/geometry-test.lua
+++ b/city/geometry-test.lua
@@ -360,3 +360,61 @@ test(
       assert(n.point.y == 0.25)
    end
 )
+
+
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- geom.quadtree
+--
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+test(
+   'use quadtree to find closest points',
+   function()
+      local tree = geom.quadtree(w, h)
+
+      local N = 1000
+
+      -- insert N random points
+      local points = {}
+      for i=1,N do
+	 local point = geom.point(math.random(), math.random())
+	 table.insert(points, point)
+	 tree:insert(point)
+      end
+
+      local find_closest = function(point)
+	 local dist = math.huge
+	 local closest
+	 for _, pt in ipairs(points) do
+	    if point:distance_to(pt) < dist then
+	       dist = point:distance_to(pt)
+	       closest = pt
+	    end
+	 end
+	 return closest
+      end
+
+      -- check closest for N random points
+      for i=1,N do
+	 local point = geom.point(math.random(), math.random())
+	 local closest_bf = find_closest(point)
+	 local closest_qt = tree:get_closest(point)
+
+	 assert(closest_bf.x == closest_qt.x,
+		string.format(
+		   'bf: (%0.3f, %0.3f), qt: (%0.3f, %0.3f)',
+		   closest_bf.x, closest_bf.y,
+		   closest_qt.x, closest_qt.y
+		)
+	 )
+	 assert(closest_bf.y == closest_qt.y,
+		string.format(
+		   'bf: (%0.3f, %0.3f), qt: (%0.3f, %0.3f)',
+		   closest_bf.x, closest_bf.y,
+		   closest_qt.x, closest_qt.y
+		)
+	 )
+      end
+   end
+)
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