BinarySearchTree.lua

------------------------------------------------------------------------------------------------------------------------------------------
-- @ CloneTrooper1019, 2021
--   Luau BinarySearchTree
------------------------------------------------------------------------------------------------------------------------------------------

local BinarySearchTree = {}
BinarySearchTree.__index = BinarySearchTree

local VALUE = 1
local LEFT_NODE = 2
local RIGHT_NODE = 3

local alloc = table.create
local push = table.insert
local pop = table.remove

function BinarySearchTree.new()
	local tree =
	{
		Size = 0;
		Root = {};
		Nodes = {};
	}
	
	return setmetatable(tree, BinarySearchTree)
end

function BinarySearchTree:Has(value)
	return (self.Nodes[value] ~= nil)
end

function BinarySearchTree:Insert(value)
	if not value then
		return false
	end

	if self.Nodes[value] then
		return true
	end
	
	local node = self.Root
	local height = 0
	
	while node[VALUE] do
		if value < node[VALUE] then
			node = node[LEFT_NODE]
		elseif value > node[VALUE] then
			node = node[RIGHT_NODE]
		else
			return false
		end
		
		height += 1
	end
	
	node[VALUE] = value
	node[LEFT_NODE] = alloc(3)
	node[RIGHT_NODE] = alloc(3)
	
	self.Nodes[value] = node
	self.Size += 1
	
	return true
end

function BinarySearchTree:Delete(oldValue)
	local nodes = self.Nodes
	local oldNode = nodes[oldValue]

	if not oldNode then
		return false
	end

	local left = oldNode[LEFT_NODE]
	local right = oldNode[RIGHT_NODE]
	
	if left[VALUE] and right[VALUE] then
		local newNode = right

		while newNode[LEFT_NODE][VALUE] do
			newNode = newNode[LEFT_NODE]
		end
		
		local newValue = newNode[VALUE]
		oldNode[VALUE] = newValue
		newNode[VALUE] = oldValue

		nodes[oldValue] = newNode
		nodes[newValue] = oldNode

		return self:Delete(oldValue)
	else
		local delete = oldNode
		local copyNode
		
		if left[VALUE] then
			copyNode = left
		elseif right[VALUE] then
			copyNode = right
		end
		
		nodes[oldValue] = nil
		
		if copyNode then
			local newValue = copyNode[VALUE]
			delete = copyNode
			
			oldNode[VALUE] = newValue
			nodes[newValue] = oldNode
			
			oldNode[LEFT_NODE] = copyNode[LEFT_NODE]
			oldNode[RIGHT_NODE] = copyNode[RIGHT_NODE]
		end

		if delete == self.Root then
			self.Root = copyNode
		end
		
		delete[VALUE] = nil
		delete[LEFT_NODE] = nil
		delete[RIGHT_NODE] = nil

		self.Size -= 1
		return true
	end
end

function BinarySearchTree:Iterate()
	return coroutine.wrap(function ()
		local visited = alloc(self.Size)
		local stack = {}
		
		local node = self.Root
		local index = 1
		
		while node do
			-- Visit the left node
			local left = node[LEFT_NODE]
			
			if left and not visited[left] then
				table.insert(stack, node)
				node = left
			else
				-- Visit this node
				if not visited[node] then
					local value = node[VALUE]
					
					if value then
						coroutine.yield(index, value)
						index += 1
					end
					
					visited[node] = true
				end
				
				-- Visit the right node
				local right = node[RIGHT_NODE]
				
				if right and not visited[right] then
					table.insert(stack, node)
					node = right
				else
					node = pop(stack)
				end
			end
		end
	end)
end

function BinarySearchTree:ToArray()
	local array = alloc(self.Size)
	
	for i, data in self:Iterate() do
		array[i] = data
	end
	
	return array
end

return BinarySearchTree