Playing with quadtrees in Lua

quadtree.lua

IMAGE1 = {
   ".#..............",
   ".#....#..#..##.#",
   ".#....#..#.#..##",
   ".####..##...##.#",
   "................",
   "..###########...",
   "..#.........###.",
   "..#.........#..#",
   "..#.........###.",
   "...#########....",
   "................",
   ".###..##...##..#",
   ".#...#..#.#..#.#",
   ".##..#..#.#..#.#",
   ".#...#..#.#..#..",
   ".#....##...##..#",
}

IMAGE2 = {
   ".#.#.#.#",
   "#.#.#.#.",
   ".#.#.#.#",
   "#.#.#.#.",
   ".#.#.#.#",
   "#.#.#.#.",
   ".#.#.#.#",
   "#.#.#.#.",
}

IMAGE3 = {
   "#......#..####..",
   ".#....#..#....#.",
   "..#..#..#......#",
   "...##...#......#",
   "...##...#......#",
   "..#..#..#......#",
   ".#....#..#....#.",
   "#......#..####..",
   "..####..#......#",
   ".#....#..#....#.",
   "#......#..#..#..",
   "#......#...##...",
   "#......#...##...",
   "#......#..#..#..",
   ".#....#..#....#.",
   "..####..#......#",
}

IMAGE4 = {
   "...##...",
   "..####..",
   ".######.",
   "########",
   "...##...",
   "...##...",
   "...##...",
   "...##..."
}

IMAGE5 = {
   "...#",
   "..##",
   ".###",
   "####"
}

IMAGE6 = {
   "........",
   "........",
   ".....#..",
   "........",
   "........",
   "........",
   "........",
   "........"
}

function get(img, x, y)
   return img[y]:sub(x,x) ~= '.'
end

function equal(subtree1, subtree2)
   if type(subtree1) == 'table' and type(subtree2) == 'table' then
      for k,v in ipairs(subtree1) do
         if not equal(v, subtree2[k]) then return false end
      end
      return true
   else
      return subtree2 == subtree1
   end
end

-- Generates and returns a tree from a size*size
-- region with its upper-left corner at (x,y).
-- Coords are 1-based with (1,1) being top left.
-- Leaves contain booleans
function generate_tree(image, x, y, size)
   local tree = {}

   if size == 1 then
      return get(image, x, y)
   else
      local ns = size / 2
      tree[1] = generate_tree(image, x, y, ns)
      tree[2] = generate_tree(image, x + ns, y, ns)
      tree[3] = generate_tree(image, x, y + ns, ns)
      tree[4] = generate_tree(image, x + ns, y + ns, ns)
   end

   return tree
end

-- Takes a tree and a list of flyweights, returns
-- an index into the flyweight list and the
-- flyweight list (after having put more tokens
-- into it)
function compress_tree(tree, flyweights)
   flyweights = flyweights or {false, true}

   if type(tree) == 'table' then
      local new_tree = {}
      for k,v in ipairs(tree) do
         new_tree[k] = compress_tree(v, flyweights)
      end
      tree = new_tree
   end

   for index, flyweight in ipairs(flyweights) do
      if equal(tree, flyweight) then
         return index, flyweight
      end
   end

   table.insert(flyweights, tree)
   return #flyweights, flyweights
end

function print_flyweight_table(flyweights)
   for index, f in ipairs(flyweights) do
      if type(f) == 'table' then
         print(index, f[1], f[2], f[3], f[4])
      else
         print(index, f)
      end
   end
end

-- Read an image into a quadtree, compress it,
-- then print out the flyweight table.
function read_and_compress(img)
   local t = generate_tree(img, 1, 1, #img)
   local i, f = compress_tree(t)
   print_flyweight_table(f)
end

read_and_compress(IMAGE1)