trie.lua

local function iterstr(str)
	return str:gmatch(".")
end

local function count(obj)
	local x = 0

	for k, v in pairs(obj) do
		x = x + 1
	end

	return x
end

local NODE = {}

NODE.children = nil
NODE.data = nil
NODE.bLeaf = false

NODE.__index = NODE

function Node()
	local obj = {}

	setmetatable(obj, NODE)
	NODE.__index = NODE

	obj.children = {}

	return obj
end

local TRIE = {}
TRIE.root = nil

TRIE.__index = TRIE

function Trie()
	local obj = {}

	setmetatable(obj, TRIE)
	TRIE.__index = TRIE

	obj.root = Node()

	return obj
end

function TRIE:__call(iterableKey)
	return self:Get(iterableKey)
end

function TRIE:Get(iterableKey)
	local current = self.root

	for x in iterableKey do
		local child = current.children[x]

		if (not child) then
			return
		end

		current = child
	end

	return current.data, current.bLeaf
end

function TRIE:Set(iterableKey, value)
	local current = self.root

	for x in iterableKey do
		local child = current.children[x]

		if (not child) then
			child = Node()

			for k, v in pairs(child.children) do
				print(k, v)
			end

			current.children[x] = child
		end

		current = child
	end

	current.bLeaf = true
	current.data = value
end

function TRIE:Remove(key, node, depth, parent)
	parent = parent or self.root
	node = node or self.root
	depth = depth or 1
	local index = key:sub(depth, depth)

	if (not node) then
		return
	end

	if (depth == #key + 1) then
		node.bLeaf = false

		if (count(node.children) <= 0) then
			for k, v in pairs(parent.children) do
				if (v == node) then
					parent.children[k] = nil

					break
				end
			end

			node = nil
		end

		return node
	end

	local ret = self:Remove(key, node.children[index], depth + 1, node)

	if (count(node.children) <= 0 and not node.bLeaf) then
		for k, v in pairs(parent.children) do
			if (v == node) then
				parent.children[k] = nil

				break
			end
		end

		node = nil
	end

	return node
end

function TRIE:GetLongestPrefix(iterableKey)
	local current = self.root
	local prefix = {}
	local value = nil
	local stack = {}

	for x in iterableKey do
		if (current.data) then
			value = current.data

			for _, v in ipairs(stack) do
				table.insert(prefix, v)
			end

			stack = {}
		end

		local child = current.children[x]

		if (not current.children[x]) then
			break
		end

		table.insert(stack, x)
		current = current.children[x]
	end

	if (current.data) then
		value = current.data

		for _, v in ipairs(stack) do
			table.insert(prefix, v)
		end
	end

	return prefix, value
end

local keys = {
	"abc",
	"def",
	"bath",
	"bathe"
}

local trie = Trie()

for k, v in ipairs(keys) do
	trie:Set(iterstr(v), k)
end

function recurse(obj, visited)
	visited = visited or {}

	if (not obj) then
		return
	end

	for letter, node in pairs(obj.children) do
		if (not visited[node]) then
			visited[node] = true

			print(letter, node.data, node.bLeaf)

			recurse(node, visited)
		end
	end
end

trie:Remove("bathe")

local text = "bathe"

local match, bLeaf = trie(iterstr(text))
print(("Match: %s - %s"):format(match and "Found" or "No match", bLeaf and "Is leaf" or "Not leaf"))

local prefix, value = trie:GetLongestPrefix(iterstr(text))

print(("Longest prefix: %s"):format(table.concat(prefix, "")))