avl tree

Avl.luau

--!strict

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

local Queue = require("./Queue")

local Avl = {}
Avl.__index = Avl

export type AvlNode<K, V> = {
	index: K,
	data: V,
	height: number,

	parent: AvlNode<K, V>?,
	right: AvlNode<K, V>?,
	left: AvlNode<K, V>?,
}

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

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

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

function Avl.get<K, V>(self: Avl<K, V>, key: K): V?
	-- Gets a value wow!
	local node = self:_query(key)
	return node and node.data or nil
end

function Avl.range<K, V>(
	self: Avl<K, V>,
	callback: (K, V) -> (),
	min: K,
	max: K
)
	type Node = AvlNode<K, V>
	
	local queue = Queue.new() :: Queue.Queue<Node>
	local compare = self.compare
	
	queue:enqueue(self.root :: Node)
	
	-- Do a BFS through the tree and only call the callback
	-- if the index is within the range. Moreover, we check
	-- within the interval (min, max). min and max are not included.
	while not queue:empty() do
		local data = queue:dequeue()
		
		local index = data.index
		local hasMin = compare(min, index)
		local hasMax = compare(index, max)
		
		if data.left and hasMin then
			queue:enqueue(data.left)	
		end
		
		if data.right and hasMax then
			queue:enqueue(data.right)
		end
		
		if hasMin and hasMax then
			callback(index, data.data)
		end
	end
end

function Avl.delete<K, V>(self: Avl<K, V>, value: K)
	local x = self:_query(value)
	if not x then
		return
	end

	-- 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.index, predecessor.index = predecessor.index, x.index

			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.index, successor.index = successor.index, x.index

			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.insert<K, V>(self: Avl<K, V>, index: K, value: V): AvlNode<K, V>?
	-- 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,
			index = index,
			height = 0,
		}

		return self.root
	end

	local compare = self.compare
	while node do
		if node.index == index then
			-- Well, we can't do anything
			return node
		elseif compare(node.index, index) then
			-- Go right
			if not node.right then
				local result = {
					height = 0,
					index = index,
					data = value,
					parent = node,
				}

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

				return result
			end

			node = node.right
		else
			-- Go left
			if not node.left then
				local result = {
					height = 0,
					index = index,
					data = value,
					parent = node,
				}

				-- Second step, rebalance the tree
				node.left = result
				self:_rebalance(result)

				return result
			end

			node = node.left
		end
	end

	return node
end

function Avl.inorder<K, V>(self: Avl<K, V>, callback: (key: K, value: V) -> ())
	-- Inorder traversal goes from left -> root -> right
	local stack = {}
	local count = 1

	local current: AvlNode<K, V>? = self.root
	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.index, current.data)
			current = current.right
		end
	end
end

function Avl.getHeight<K, V>(self: Avl<K, V>, x: AvlNode<K, V>?): 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<K, V>(self: Avl<K, V>, x: AvlNode<K, V>): AvlNode<K, V>
	-- Find left most node
	while x.left do
		x = x.left
	end

	return x
end

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

	return x
end

function Avl.getLastInorder<K, V>(self: Avl<K, V>, x: AvlNode<K, V>): AvlNode<K, V>
	-- 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<K, V>(self: Avl<K, V>, x: AvlNode<K, V>): AvlNode<K, V>
	-- 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<K, V>(self: Avl<K, V>, x: AvlNode<K, V>): AvlNode<K, V>?	
	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<K, V>(self: Avl<K, V>, x: AvlNode<K, V>): AvlNode<K, V>?	
	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._rotateLeft<K, V>(self: Avl<K, V>, x: AvlNode<K, V>)
	--   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.index, y.index = y.index, x.index
	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._rotateRight<K, V>(self: Avl<K, V>, x: AvlNode<K, V>)
	--       x          y
	--   y      C => A     x
	-- A   B             B   C
	local y = assert(x.left)
	local B = y.right
	local A = y.left

	-- Swap
	x, y = y, x
	x.index, y.index = y.index, x.index
	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<K, V>(self: Avl<K, V>, a: AvlNode<K, V>?)
	-- 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._query<K, V>(self: Avl<K, V>, index: K): AvlNode<K, V>?
	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 do
		if node.index == index then
			break
		elseif compare(node.index, index) then
			-- Go right
			node = node.right
		else
			-- Go left
			node = node.left
		end
	end

	if not node or node.index ~= index then
		-- Not the same
		return nil
	end

	return node
end

return Avl

Queue.luau

--!strict

-- yet another queue implementation :P
-- magicoal_nerb

local Queue = {}
Queue.__index = Queue

export type Queue<T> = typeof(setmetatable({} :: {
	data: { T },
	size: number,
}, Queue))

function Queue.new<T>(): Queue<T>
	return setmetatable({
		data = {},
		size = 0,
	}, Queue)
end

function Queue.empty<T>(self: Queue<T>): boolean
	return self.size == 0
end

function Queue.iterate<T>(self: Queue<T>, callback: (T, ...any) -> (), ...: any)
	for i, object in self.data do
		callback(object, ...)
	end
end

function Queue.enqueue<T>(self: Queue<T>, data: T)
	self.data[self.size + 1] = data
	self.size += 1
end

function Queue.dequeue<T>(self: Queue<T>): T
	self.size -= 1
	return self.data[self.size + 1]
end

function Queue.clear<T>(self: Queue<T>)
	self.size = 0
	table.clear(self.data)
end

return Queue