avltree.lua

local function explode(str, separator)
	local out = {}

	for w in str:gmatch("([^" .. separator .. "]+)") do out[#out + 1] = w end

	return out
end

local NODE	= {}

NODE.key	= 0
NODE.value	= ""
NODE.left	= nil
NODE.right	= nil
NODE.height	= 1

function Node(key, data)
	local obj	= {}
	obj.key		= key
	obj.value	= data or ""

	setmetatable(obj, NODE)
	NODE.__index = NODE

	return obj
end

local AVLTREE	= {}

AVLTREE.root	= nil

local function MaxHeight(a, b)
	return math.max(a and a.height or 0, b and b.height or 0)
end

local function SmallestKey(node)
	while (node and node.left) do
		node = node.left
	end

	return node
end

function AvlTree(data)
	local obj = {}

	setmetatable(obj, AVLTREE)
	AVLTREE.__index = AVLTREE

	return obj
end

function AVLTREE:GetBalance(node)
	return node and (node.left and node.left.height or 0) - (node.right and node.right.height or 0) or 0
end

function AVLTREE:ShiftRight(node)
	local left = node.left
	local right = left.right

	left.right = node
	node.left = right

	node.height = MaxHeight(node.left, node.right) + 1
	left.height = MaxHeight(left.left, left.right) + 1

	return left
end

function AVLTREE:ShiftLeft(node)
	local right = node.right
	local left = right.left

	right.left = node
	node.right = left

	node.height	 = MaxHeight(node.left, node.right) + 1
	right.height = MaxHeight(right.left, right.right) + 1

	return right
end

function AVLTREE:Insert_(key, value, node)
	if (not node) then
		return Node(key, value)
	end

	if (key < node.key) then
		node.left = self:Insert_(key, value, node.left)
	elseif (key > node.key) then
		node.right = self:Insert_(key, value, node.right)
	else
		return node
	end

	node.height = MaxHeight(node.left, node.right) + 1

	local balance = self:GetBalance(node)

	if (balance > 1 and key < node.left.key) then
		return self:ShiftRight(node)
	elseif (balance < -1 and key > node.right.key) then
		return self:ShiftLeft(node)
	elseif (balance > 1 and key > node.left.key) then
		node.left = self:ShiftLeft(node.left)

		return self:ShiftRight(node)
	elseif (balance < -1 and key < node.right.key) then
		node.right = self:ShiftRight(node.right)

		return self:ShiftLeft(node)
	end

	return node
end

function AVLTREE:Insert(key, value)
	self.root = self:Insert_(key, value, self.root)
end

function AVLTREE:Get(key)
	local node = self.root

	while (node) do
		if (node.key == key) then
			return node.value
		elseif (node.key < key) then
			node = node.right
		elseif (node.key > key) then
			node = node.left
		end
	end

	return "Not Found"
end

function AVLTREE:Remove_(key, node)
	if (not node) then
		return node
	end

	if (key < node.key) then
		node.left = self:Remove_(key, node.left)
	elseif (key > node.key) then
		node.right = self:Remove_(key, node.right)
	else
		if (not node.left or not node.right) then
			local temp = node.left or node.right

			return temp
		end

		local temp = SmallestKey(node.right)

		node.key = temp.key
		node.value = temp.value
		node.right = self:Remove_(temp.key, node.right)
	end

	node.height = MaxHeight(node.left, node.right) + 1

	local balance = self:GetBalance(node)

	if (balance > 1 and self:GetBalance(node.left) >= 0) then
		return self:ShiftRight(node)
	elseif (balance > 1 and self:GetBalance(node.left) < 0) then
		node.left = self:ShiftLeft(node.left)

		return self:ShiftRight(node)
	elseif (balance < -1 and self:GetBalance(node.right) <= 0) then
		return self:ShiftLeft(node)
	elseif (balance < -1 and self:GetBalance(node.right) > 0) then
		node.right = self:ShiftRight(node.right)

		return self:ShiftLeft(node)
	end

	return node
end

function AVLTREE:Remove(key)
	self.root = self:Remove_(key, self.root)
end

function AVLTREE:LevelOrder()
	local height = self.root.height

	local function printLevel(node, level)
		if (not node) then
			return
		end

		if (level == 1) then
			io.stdout:write(("%s:%s(%s) "):format(tostring(node.key), tostring(node.value), tostring(self:GetBalance(node))))
		elseif (level > 1) then
			printLevel(node.left, level - 1)
			printLevel(node.right, level - 1)
		end
	end

	for i = 1, height do
		printLevel(self.root, i)
	end

	io.stdout:write("\n")
end

-- local tree = AvlTree()

-- local commands = {
-- 	put			= function(args) tree:Insert(tonumber(args[2]), table.concat(args, " ", 3))	end,
-- 	get			= function(args) io.stdout:write(tree:Get(tonumber(args[2])) .. "\n")		end,
-- 	remove		= function(args) tree:Remove(tonumber(args[2]))								end,
-- 	levelorder	= function(args) tree:LevelOrder()											end
-- }

-- while (text ~= "") do
-- 	local text = io.stdin:read()

-- 	if (not text or text == "") then break end

-- 	local packed = explode(text, " ")

-- 	if (not tonumber(packed[1])) then
-- 		local command = string.lower(packed[1])

-- 		commands[command](packed)
-- 	end
-- end