bk-tree.lua

---------------------------
-- bk-tree datastructure
--
-- http://en.wikipedia.org/wiki/BK-tree
-- @module bk-tree
-- @author Robin Hübner
-- robinhubner@gmail.com
-- @release version 1.0.3
-- @license MIT

local bk_tree = {}


local function lazy_copy(t1)

	local cp = {}

	for k, v in pairs(t1) do
		cp[k] = v
	end

	return cp

end

local function min(a, b, c)

	local min_val = a

	if b < min_val then min_val = b end
	if c < min_val then min_val = c end

	return min_val

end

----------------------------------
--- Levenshtein distance function.
-- @tparam string s1
-- @tparam string s2
-- @treturn number the levenshtein distance
-- @within Metrics
function bk_tree.levenshtein_dist(s1, s2)
	
	if s1 == s2 then return 0 end
	if s1:len() == 0 then return s2:len() end
	if s2:len() == 0 then return s1:len() end
	if s1:len() < s2:len() then s1, s2 = s2, s1 end

	local t = {}
	for i=1, #s1+1 do
		t[i] = {i-1}
	end

	for i=1, #s2+1 do
		t[1][i] = i-1
	end

	local cost
	for i=2, #s1+1 do
	
		for j=2, #s2+1 do
			cost = (s1:sub(i-1,i-1) == s2:sub(j-1,j-1) and 0) or 1
			t[i][j] = min(
				t[i-1][j] + 1,
				t[i][j-1] + 1,
				t[i-1][j-1] + cost)
		end

	end

	return t[#s1+1][#s2+1]
	
end

function bk_tree.hook(param)

	local name, callee = debug.getlocal(2, 1)
	local f = debug.getinfo(2, "f").func
	local p = debug.getinfo(3, "f").func
	--[[ previous function in the callstack, if called from the same place,
			don't add to the insert/remove counters. ]]--

	if f == bk_tree.insert and p ~= bk_tree.insert then
		callee.stats.nodes = callee.stats.nodes + 1
	elseif f == bk_tree.remove and p ~= bk_tree.remove then
		callee.stats.nodes = callee.stats.nodes - 1
	elseif f == bk_tree.query and p == bk_tree.query then
		callee.stats.queries = callee.stats.queries + 1
	end

end

--- Hooks debugging into tree execution.
-- Keeps track of number of nodes created, queries made,
-- note that this must be run directly after tree is created
-- in order to get correct information.
-- @within Debug
--- @usage
-- local bktree = require "bk-tree"
-- local tree = bktree:new("word")
-- tree:debug()
-- tree:insert("perceive")
-- tree:insert("beautiful")
-- tree:insert("definitely")
-- local result = tree:query("definately", 3)
-- tree:print_stats()
--
-- -- output
-- Nodes: 4
-- Queries: 3
-- Nodes Queried: 75%
function bk_tree:debug()

	local nc = 0
	if self.root then nc = 1 end
	self.stats = { nodes = nc, queries = 0 }
	debug.sethook(self.hook, "c")

end

--- Print execution stats.
-- Prints nodes queried and total nodes, as well as a fraction of 
-- nodes visited to satisfy the query, resets the counter of nodes queried when called.
-- @within Debug
-- @see debug
function bk_tree:print_stats()

	print("\nNodes: " .. self.stats.nodes)
	print("Queries: " .. self.stats.queries)
	print("Nodes Queried: " .. self.stats.queries/self.stats.nodes*100 .. "%\n")
	self.stats.queries = 0

end

--- Fetch execution stats.
-- Returns a copy of the execution stats that @{print_stats} would print, requires debug to have been enabled
-- to not just return defaults. Useful if you want to profile things.
-- @within Debug
-- @return {key = value,...}
function bk_tree:get_stats()

	return lazy_copy(self.stats)

end

---------------------------
--- Creates a new bk-tree.
-- @constructor
-- @string[opt] root_word the root of the new tree
-- @tparam[opt=levenshtein_dist] function dist_func the distance function used
-- @see levenshtein_dist
-- @return the new bk-tree instance
--- @usage
-- local bktree = require "bk-tree"
-- local tree = bktree:new("word")
function bk_tree:new(root_word, dist_func)

	local n_obj = {}
	if root_word then n_obj.root = { str = root_word, children = {} } end
	n_obj.dist_func = dist_func or self.levenshtein_dist

	setmetatable(n_obj, self)
	self.__index = self

	return n_obj

end

--------------------------------
--- Inserts word into the tree.
-- @string word
-- @treturn bool true if inserted, false if word already exists in tree
--- @usage
-- local bktree = require "bk-tree"
-- local tree = bktree:new("root")
-- local success = tree:insert("other_word")
function bk_tree:insert(word, node)

	node = node or self.root

	if not node then
		self.root = { str = word, children = {} }
		return true
	end	

	local dist = self.dist_func(word, node.str)
	if dist == 0 then return false end

	local some_node = node.children[dist]

	if not some_node then
		node.children[dist] = { str = word, children = {} }
		return true
	end	

	return self:insert(word, some_node)

end

--------------------------------
--- Query the tree for matches.
-- @string word
-- @tparam number n max edit distance to use when querying
-- @treturn {{str=string,distance=number},....} table of tables with matching words, empty table if no matches
--- @usage
-- local bktree = require "bk-tree"
-- local tree = bktree:new("word")
-- tree:insert("hello")
-- tree:insert("goodbye")
-- tree:insert("woop")
-- local result = tree:query("woop", 1)
function bk_tree:query(word, n, node, matches)

	node = node or self.root
	matches = matches or {}

	if not node then return matches end

	local dist = self.dist_func(word, node.str)
	if dist <= n then matches[#matches+1] = {str = node.str, distance = dist} end
	
	for k, child in pairs(node.children) do
		if child ~= nil then
			if k >= dist-n and k <= dist+n then
				self:query(word, n, child, matches)
			end
		end
	end

	return matches

end

---------------------------------------------------------
--- Queries the the tree for a match, sorts the results.
-- Calls @{query} and returns the results sorted.
-- @string word
-- @tparam number n max edit distance to use when querying
-- @treturn {{str=string,distance=number},....} table of tables with matching words sorted by distance, empty table if no matches
--- @usage
-- local bktree = require "bk-tree"
-- local tree = bktree:new("word")
-- tree:insert("woop")
-- tree:insert("worp")
-- tree:insert("warp")
-- local result = tree:query_sorted("woop", 3)
function bk_tree:query_sorted(word, n)

	local result = self:query(word, n)

	table.sort(result, function(a,b) return a.distance < b.distance end)

	return result

end

return bk_tree