avl tree

Avl.luau

--!strict

-- AVL Tree implementations based on the mit ocw lecture videos
-- poopbarrel/magicoal_nerb :^)

local Avl = {}
Avl.__index = Avl

export type AvlNode<T> = {
	data: T,
	height: number,
	
	parent: AvlNode<T>?,
	right: AvlNode<T>?,
	left: AvlNode<T>?,
}

export type Avl<T> = typeof(setmetatable({} :: {
	compare: (a: T, b: T) -> boolean,
	root: AvlNode<T>?,
}, Avl))

local function height<T>(node: AvlNode<T>?)
	-- Helper that cleans up getting the height
	-- of a node
	if node then
		return node.height
	else
		return -1
	end
end

function Avl.new<T>(compare: (a: T, b: T) -> boolean): Avl<T>
	return setmetatable({
		compare = compare,
		root = nil,
	}, Avl)
end

function Avl.getHeight<T>(self: Avl<T>, x: AvlNode<T>?): number
	if not x then
		-- Base case so we get 1 - 1 = 0 for leaf nodes
		return -1
	end
	
	return 1 + math.max(self:getHeight(x.left), self:getHeight(x.right))
end

function Avl.findMin<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>
	-- Find left most node
	while x.left do
		x = x.left
	end
	
	return x
end

function Avl.findMax<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>
	-- Find right most node
	while x.right do
		x = x.right
	end
	
	return x
end

function Avl.getLastInorder<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>
	-- Gets the last inorder element of the given node
	while x and x.right do
		x = x.right
	end

	return x
end

function Avl.getFirstInorder<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>
	-- Gets the first inorder element of the given node
	while x and x.left do
		x = x.left
	end
	
	return x
end

function Avl.predecessor<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>?	
	if x.left then
		-- If a left node exists, then our predecessor is down the tree
		return self:getLastInorder(x.left)
	end
	
	-- If the current node is to the right of its parent, then that parent
	-- must be the predecessor
	while x.parent and x.parent.left == x do
		x = x.parent
	end
	
	return x.parent
end

function Avl.successor<T>(self: Avl<T>, x: AvlNode<T>): AvlNode<T>?	
	if x.right then
		-- If a right node exists, then our successor is down the tree
		return self:getFirstInorder(x.right)
	end
	
	-- If the current node is to the left of its parent, then that parent
	-- must be the succesor
	while x.parent and x.parent.right == x do
		x = x.parent
	end
	
	return x.parent
end

function Avl.delete<T>(self: Avl<T>, x: AvlNode<T>?)
	assert(x)
	
	-- First step, do an ordinary BST deletion
	-- It's a branch, replace it with the predecessor
	while x.left or x.right do
		if x.left then
			-- If we have a left child, then we swap with the predecessor
			-- somewhere inside our tree
			local predecessor = self:predecessor(x)
			assert(predecessor)
			x.data, predecessor.data = predecessor.data, x.data
			x = predecessor
		else
			-- If we have a right child, then we swap with the successor
			-- somewhere inside our tree
			local successor = self:successor(x)
			assert(successor)
			x.data, successor.data = successor.data, x.data
			x = successor
		end
	end
	
	local parent = x.parent
	if parent then
		-- Remove and rebalance the tree
		if parent.left == x then
			parent.left = nil
		else
			parent.right = nil
		end
		
		self:rebalance(parent)
	end
end

function Avl.rotateLeft<T>(self: Avl<T>, x: AvlNode<T>)
	--   x                 y
	-- A     y    =>   x     C
	--     B   C     A   B
	local y = assert(x.right)
	local A = x.left
	local B = y.left
	local C = y.right
	
	-- Swap
	x, y = y, x
	x.data, y.data = y.data, x.data
	y.left = x
	y.right = C
	x.left = A
	x.right = B
	
	if A then
		A.parent = x
	end
	
	if C then
		C.parent = y
	end	
end

function Avl.inorder<T>(self: Avl<T>, x: AvlNode<T>, callback: (value: T) -> ())
	-- Inorder traversal goes from left -> root -> right
	local stack = {}
	local count = 1
	
	local current: AvlNode<T>? = x
	while current or count > 0 do
		while current do
			-- Keep on going left
			table.insert(stack, current)
			count += 1
			current = current.left
		end
		
		current = table.remove(stack)
		count -= 1
		
		-- We can proceed to the right if its
		-- possible
		if current then
			callback(current.data)
			current = current.right
		end
	end
end

function Avl.rotateRight<T>(self: Avl<T>, x: AvlNode<T>)
	--       x          y
	--   y      C => A     x
	-- A   B             B   C
	local y = assert(x.left)
	local C = x.right
	local B = y.right
	local A = y.left
	
	-- Swap
	x, y = y, x
	x.data, y.data = y.data, x.data
	
	y.right = x
	y.left = A
	x.left = B
	
	if A then
		A.parent = y
	end
	
	if B then
		B.parent = x
	end
end

function Avl.rebalance<T>(self: Avl<T>, a: AvlNode<T>?)
	-- skew(a) = height(a.right) - height(a.left)
	-- height(a) = max(height(a.left), height(a.right)) + 1
	while a do
		local skew = height(a.right) - height(a.left)
		if skew == 2 then
			-- Double right heavy
			local right = assert(a.right)
			local rightSkew = height(right.right) - height(right.left)
			
			if rightSkew < 0 then
				self:rotateRight(right)
			end
			
			self:rotateLeft(a)
		elseif skew == -2 then
			-- Double left heavy
			local left = assert(a.left)
			local leftSkew = height(left.right) - height(left.left)
			
			if leftSkew > 0 then
				self:rotateLeft(left)
			end
			
			self:rotateRight(a)
		end
		
		-- Finalize height
		a.height = 1 + math.max(height(a.right), height(a.left))
		a = a.parent
	end
end

function Avl.get<T>(self: Avl<T>, value: T): AvlNode<T>?
	local node = self.root
	if not node then
		-- The tree is empty
		return nil
	end
	
	-- Does a standard binary search over our data
	local compare = self.compare
	while node.left or node.right do
		if compare(node.data, value) then
			-- Go right
			node = node.right
		else
			-- Go left
			node = node.left
		end
	end
	
	return node
end

function Avl.insert<T>(self: Avl<T>, value: T)
	-- First step, do an ordinary BST insert
	local node = self.root
	if not node then
		-- If no root node exists, then
		-- we make our value the root node
		self.root = {
			data = value,
			height = 0,
		}
		
		return
	end

	local compare = self.compare
	while node do
		if compare(node.data, value) then
			-- Go right
			if not node.right then
				local result = {
					height = 0,
					data = value,
					parent = node,
				}

				-- Second step, rebalance the tree
				node.right = result
				self:rebalance(result)
				return
			end

			node = node.right
		else
			-- Go left
			if not node.left then
				local result = {
					height = 0,
					data = value,
					parent = node,
				}
				
				-- Second step, rebalance the tree
				node.left = result
				self:rebalance(result)
				return
			end
			
			node = node.left
		end
	end
end

return Avl