Commonly used tree data structures in Lua

Resources Community Resources
#1

I don’t really use this and its unlikely that I ever will so I am releasing it here. This thread contains some data structures that are commonly used outside of Roblox.

The code is not perfectly written because I wrote a little more than a year ago. Hopefully somebody finds it useful.

Binary search tree

--[[
	BST
	
	not auto balance
	
--]]

-- Constants

-- efficieny variables
local setmetatable=setmetatable

-- local variables

--[[
	Tree: Root
	
	Node: Parent, Left, Right, Data
	
--]]

local function newnode(data)
	return{
		false,-- Parent
		false,-- Left
		false,-- Right
		data,}
end

local function find(tree,data)
	local c=tree.root
	while c do
		local d=c[4]
		if data<d then
			c=c[2]
		elseif data>d then
			c=c[3]
		else
			return c
		end
	end
end

local function lowerbound(tree,v)
	-- returns node if exists, else nil
	local y
	local x=tree.root
	
	while x do
		if x[4]>=v then
			y=x
			x=x[2]
		else
			x=x[3]
		end
	end
	
	return y
end

local function upperbound(tree,v)
	-- returns node if exists, else nil
	local y
	local x=tree.root
	
	while x do
		if x[4]>v then
			x=x[2]
		else
			y=x
			x=x[3]
		end
	end
	
	return y
end

local function insert(tree,data)
	local n=newnode(data)
	local c=tree.root
	if c then
		local l=c[2]
		local r=c[3]
		
		while true do
			if data<c[4]then
				if l then
					c=l
					r=l[3]
					l=l[2]
				else
					c[2]=n
					n[1]=c
					break
				end
			else
				if r then
					c=r
					l=r[2]
					r=r[3]
				else
					c[3]=n
					n[1]=c
					break
				end
			end
		end
	else
		tree.root=n
	end
	
	return n
end

local function deletenode(tree,n)
	-- find successor
	local right=n[3]
	local left=n[2]
	if right and left then
		local succ
		local temp=right[2]
		while temp do
			succ=temp
			temp=succ[2]
		end
		
		local par=n[1]
		
		if succ then
			if par then
				if par[2]==n then
					par[2]=succ
				else
					par[3]=succ
				end
			end
			
			temp=succ[3]
			succ[1][2]=temp
			succ[3]=right
			
			if temp then
				temp[1]=succ[1]
			end
			right=succ
		else
			if par then
				if par[2]==n then
					par[2]=right
				else
					par[3]=right
				end
			end
		end
		right[1]=par
		right[2]=left
	elseif right then
		local par=n[1]
		if par then
			if par[2]==n then
				par[2]=right
			else
				par[3]=right
			end
		end
		right[1]=par
	elseif left then
		right=left
		local par=n[1]
		if par then
			if par[2]==n then
				par[2]=right
			else
				par[3]=right
			end
		end
		right[1]=par
	end
	
	if n==tree.root then
		tree.root=right
	end
end

local function deletedata(tree,data)
	deletenode(tree,find(tree,data))
end

local function empty(tree)
	return tree.root==nil
end

local function printnode(node,prefix,istail)
	if not node then
		return''
	end
	
	local s='\n'..prefix..((istail and'\226\148\148\226\148\128\226\148\128 ')or'\226\148\156\226\148\128\226\148\128 ')..node[4]
	prefix=prefix..((istail and'    ')or'\226\148\130   ')
	
	local l=node[2]
	local r=node[3]
	
	if l and r then
		-- right then left first so if you do 90º clockwise Rotation it looks like a BST
		return s..printnode(r,prefix)..printnode(l,prefix,true)
	elseif l then
		return s..printnode(l,prefix,true)
	elseif r then
		return s..printnode(r,prefix,true)
	else
		return s
	end
end

local function dc(a,s,e)
	-- divide and conquer to construct a perfect BST
	if e<=s then
		e=newnode(a[s])
		e[5]=true
		return e
	else
		local m=(s+e)/2
		m=m-m%1
		
		local x=newnode(a[m])
		if s<m then
			-- edge case which causes some nodes to be inserted twice
			x[2]=dc(a,s,m-1)
			x[2][1]=x
		end
		
		x[3]=dc(a,m+1,e)
		x[3][1]=x
		
		return x
	end
end

--[[
	local function isBST(node, minKey, maxKey)
		if node == b.Nil then
			return true
		end
		
		if node[4] < minKey or node[4] > maxKey then
			return false
		end
		
		return isBST(node[2],minKey,node[4]-1) and isBST(node[3],node[4]+1,maxKey)
	end
	
	print(isBST(b.Root,-math.huge,math.huge))
--]]

local bst={
	find=find,
	insert=insert,
	deletenode=deletenode,
	deletedata=deletedata,
	empty=empty,
	
	__tostring=function(tree)
		return printnode(tree.root,'',true)
	end,
}

function bst.new(a)
	local root
	local Size
	
	if a then
		-- assumes SORTED array, if not sorted then call table.sort BEFORE passing it
		Size=#a
		root=(Size>0 and dc(a,1,Size))or nil
	else
		Size=0
	end
	
	local self={
		root=root,
		Size=Size,
		
		empty=empty,
		insert=insert,
		deletenode=deletenode,
		deletedata=deletedata,
		find=find,
		lowerbound=lowerbound,
		upperbound=upperbound,
	}
	
	setmetatable(self,bst)
	return self
end

return bst
15 Likes
#2

Priority queue-Pairing heap

--[[
	Pairing Min Heap
	
	Binary Heap is faster for increasing and random pushes but this is faster for joining (merging) and decreasing pushes
	It is very optimized
	
	Use Kill to GC after finished using
	
--]]

--[[
	Node: {
	Key,
	LeftChild,
	NextSibling,
	}
	
--]]

-- Constants

-- cached variables

-- local variables

local tables={}
local tSize=0

local function merge(a,b)
	if not a then
		return b
	elseif not b then
		return a
	elseif a[1]<b[1]then
		-- inlined add child
		b[3]=a[2]
		a[2]=b
		
		return a
	else
		-- inlined add child
		a[3]=b[2]
		b[2]=a
		
		return b
	end
end

local function twopassmerge(a)
	local b=a and a[3]
	if b then
		local newnode=b[3]
		
		a[3]=false
		b[3]=false
		
		return merge(merge(a,b),twopassmerge(newnode))
	else
		return a
	end
end

local function empty(heap)
	return not heap.root
end

local function top(heap)
	return heap.root[1]
end

local function push(heap,key)
	local root=heap.Size+1
	heap.Size=root
	local t=heap.maxSize
	heap.maxSize=(root>t and root)or t
	
	-- inlined merge (optimized for insert)
	root=heap.root
	if not root then
		if tSize>0 then
			t=tables[tSize]
			tables[tSize]=nil
			tSize=tSize-1
			t[1]=key
			heap.root=t
		else
			heap.root={
				key,
				false,
				false,
			}
		end
	elseif root[1]<key then
		if tSize>0 then
			t=tables[tSize]
			tables[tSize]=nil
			tSize=tSize-1
			t[1]=key
			t[3]=root[2]
			root[2]=t
		else
			root[2]={
				key,
				false,
				root[2],
			}
		end
	else
		if tSize>0 then
			t=tables[tSize]
			tables[tSize]=nil
			tSize=tSize-1
			t[1]=key
			t[2]=root
			heap.root=t
		else
			heap.root={
				key,
				root,
				false,
			}
		end
	end
end

local function pop(heap)
	heap.Size=heap.Size-1
	local root=heap.root
	heap.root=twopassmerge(root[2])
	root[1]=false
	root[2]=false
	root[3]=false
	tSize=tSize+1
	tables[tSize]=root
end

local function join(heap1,heap2)
	-- assumes non empty heaps
	-- don't kill heap2 after merging
	
	local a=heap1.Size+heap2.Size
	heap1.Size=a
	local b=heap1.maxSize
	heap1.maxSize=(a>b and a)or b
	
	a=heap1.root
	b=heap2.root
	
	-- inlined merge
	if a[1]<b[1]then
		-- inlined add child
		b[3]=a[2]
		a[2]=b
	else
		-- inlined add child
		a[3]=b[2]
		b[2]=a
		
		heap1.root=b
	end
	
	return heap1
end

local function kill(heap)
	local nt=tSize-heap.maxSize+heap.Size
	for i=nt+1,tSize do
		tables[i]=nil
	end
	tSize=nt
end

local heap={
	empty=empty,
	top=top,
	push=push,
	pop=pop,
	join=join,
	kill=kill,
}

function heap.new(expectedspace)
	expectedspace=expectedspace or 0
	local nt=tSize+expectedspace
	for i=tSize+1,nt do
		tables[i]={false,false,false,}
	end
	tSize=nt
	
	return{
		empty=empty,
		top=top,
		push=push,
		pop=pop,
		join=join,
		kill=kill,
		
		Size=0,
		maxSize=expectedspace,
	}
end

return heap
2 Likes
#3

Priority queue-skewheap

--[[
	Skew Min heap
	
	Not very optimized and slower than both Binary and Pairing heaps anyways
--]]

-- Constants

-- cached variables

-- local variables

local function empty(self)
	return not self[3]
end

local function merge(self,a,b)
	local head=self
	
	if not b then
		head[3]=a
	else
		while a do
			if a[1]<b[1]then-- a is less (or higher priority) than b
				head[3]=a-- the lesser tree goes on the left
				head=a-- and we work on that in the next round
				a=head[4]-- by merging its right side with b
				head[4]=head[3]-- and we move the left side to the right
			else
				head[3]=b
				head=b
				local t=a
				a=head[4]
				b=t
				head[4]=head[3]
			end
		end
		head[3]=b
	end
end

local function push(self,k,v)
	merge(self,{
		k,-- key (what its sorted by)
		v,},self[3])
end

local function pop(self)
	local result=self[3]
	merge(self,result[3],result[4])
	
	return result[1],result[2]
end

local function peak(self)
	local result=self[3]
	
	return result[1],result[2]
end

local heap={
	empty=empty,
	merge=merge,
	push=push,
	pop=pop,
	peak=peak,
}

function heap.new()
	return{
		empty=empty,
		merge=merge,
		push=push,
		pop=pop,
		peak=peak,
	}
end

return heap
3 Likes
#4

Balanced BST-splay

--[[
	Splay Tree
	
	DONT USE
	Use RBT over this, this is SLOW
	
	self balancing
	
	seems to be slower than RBT when inserting and searching for random elements (inserting = 4, searching = 1.15)
	sequential insertion = 0.25 while sequential search = 1.5
	searching for same element = 0.65
	
	searching for same 100 elements = 2.5
	same 10 = 1.8
	
--]]

-- Constants

-- efficieny variables
local setmetatable=setmetatable

-- local variables

--[[
	Tree: Root
	
	Node: Parent, Left, Right, Data
	
--]]

local function newnode(data)
	return{
		false,-- Parent
		false,-- Left
		false,-- Right
		data,}
end

local function leftrotate(tree,x)
	local y=x[3]-- set y
	
	-- turn y's left subtree into x's right subtree
	local l=y[2]
	x[3]=l
	if l then
		l[1]=x
	end
	
	-- link x's Parent to y
	l=x[1]
	y[1]=l
	if not l then
		tree.root=y
	elseif x==l[2]then
		l[2]=y
	else
		l[3]=y
	end
	
	-- put x on y's left
	y[2]=x
	x[1]=y
end

local function rightrotate(tree,x)
	local y=x[2]-- set y
	
	-- turn y's right subtree into x's left subtree
	local l=y[3]
	x[2]=l
	if l then
		l[1]=x
	end
	
	-- link x's Parent to y
	l=x[1]
	y[1]=l
	if not l then
		tree.root=y
	elseif x==l[2]then
		l[2]=y
	else
		l[3]=y
	end
	
	-- put x on y's right
	y[3]=x
	x[1]=y
end

local function splay(tree,x)
	-- limit on max number
	local p=x[1]
	while p do
		local gp=p[1]
		if not gp then
			if p[2]==x then
				rightrotate(tree,p)
			else
				leftrotate(tree,p)
			end
		elseif p[2]==x and gp[2]==p then
			rightrotate(tree,gp)
			rightrotate(tree,p)
		elseif p[3]==x and gp[3]==p then
			leftrotate(tree,gp)
			leftrotate(tree,p)
		elseif p[2]==x and gp[3]==p then
			rightrotate(tree,p)
			leftrotate(tree,x[1])
		else
			leftrotate(tree,p)
			rightrotate(tree,x[1])
		end
		
		p=x[1]
	end
end

local function find(tree,data)
	local c=tree.root
	while c do
		local d=c[4]
		if data<d then
			c=c[2]
		elseif data>d then
			c=c[3]
		else
			splay(tree,c)
			return c
		end
	end
end

local function lowerbound(tree,v)
	-- returns node if exists, else nil
	local y
	local x=tree.root
	
	while x do
		if x[4]>=v then
			y=x
			x=x[2]
		else
			x=x[3]
		end
	end
	
	if y then
		splay(tree,y)
		return y
	end
end

local function upperbound(tree,v)
	-- returns node if exists, else nil
	local y
	local x=tree.root
	
	while x do
		if x[4]>v then
			x=x[2]
		else
			y=x
			x=x[3]
		end
	end
	
	if y then
		splay(tree,y)
		return y
	end
end

local function insert(tree,data)
	local n=newnode(data)
	local c=tree.root
	if c then
		local l=c[2]
		local r=c[3]
		
		while true do
			if data<c[4]then
				if l then
					c=l
					r=l[3]
					l=l[2]
				else
					c[2]=n
					n[1]=c
					break
				end
			else
				if r then
					c=r
					l=r[2]
					r=r[3]
				else
					c[3]=n
					n[1]=c
					break
				end
			end
		end
		
		splay(tree,n)
	else
		-- already root so no splay
		tree.root=n
	end
	
	tree.Size=tree.Size+1
	return n
end

local function deletenode(tree,n)
	-- find successor
	local right=n[3]
	local left=n[2]
	if right and left then
		local succ
		local temp=right[2]
		while temp do
			succ=temp
			temp=succ[2]
		end
		
		local par=n[1]
		
		if succ then
			if par then
				if par[2]==n then
					par[2]=succ
				else
					par[3]=succ
				end
			end
			
			temp=succ[3]
			succ[1][2]=temp
			succ[3]=right
			
			if temp then
				temp[1]=succ[1]
			end
			right=succ
		else
			if par then
				if par[2]==n then
					par[2]=right
				else
					par[3]=right
				end
			end
		end
		right[1]=par
		right[2]=left
		
		if par then
			splay(tree,par)
		end
	elseif right then
		local par=n[1]
		if par then
			if par[2]==n then
				par[2]=right
			else
				par[3]=right
			end
		end
		right[1]=par
		
		if par then
			splay(tree,par)
		end
	elseif left then
		right=left
		local par=n[1]
		if par then
			if par[2]==n then
				par[2]=right
			else
				par[3]=right
			end
		end
		right[1]=par
		
		if par then
			splay(tree,par)
		end
	end
	
	tree.Size=tree.Size-1
	if n==tree.root then
		tree.root=right
	end
end

local function deletedata(tree,data)
	deletenode(tree,find(tree,data))
end

local function empty(tree)
	return tree.root==nil
end

local function printnode(node,prefix,istail)
	if not node then
		return''
	end
	
	local s='\n'..prefix..((istail and'\226\148\148\226\148\128\226\148\128 ')or'\226\148\156\226\148\128\226\148\128 ')..node[4]
	prefix=prefix..((istail and'    ')or'\226\148\130   ')
	
	local l=node[2]
	local r=node[3]
	
	if l and r then
		-- right then left first so if you do 90º clockwise Rotation it looks like a BST
		return s..printnode(r,prefix)..printnode(l,prefix,true)
	elseif l then
		return s..printnode(l,prefix,true)
	elseif r then
		return s..printnode(r,prefix,true)
	else
		return s
	end
end

local function dc(a,s,e)
	-- divide and conquer to construct a perfect BST
	if e<=s then
		e=newnode(a[s])
		e[5]=true
		return e
	else
		local m=(s+e)/2
		m=m-m%1
		
		local x=newnode(a[m])
		if s<m then
			-- edge case which causes some nodes to be inserted twice
			x[2]=dc(a,s,m-1)
			x[2][1]=x
		end
		
		x[3]=dc(a,m+1,e)
		x[3][1]=x
		
		return x
	end
end

--[[
	local function isBST(node, minKey, maxKey)
		if node == b.Nil then
			return true
		end
		
		if node[4] < minKey or node[4] > maxKey then
			return false
		end
		
		return isBST(node[2],minKey,node[4]-1) and isBST(node[3],node[4]+1,maxKey)
	end
	
	print(isBST(b.Root,-math.huge,math.huge))
--]]

local splay={
	empty=empty,
	insert=insert,
	deletenode=deletenode,
	deletedata=deletedata,
	find=find,
	lowerbound=lowerbound,
	upperbound=upperbound,
	
	__tostring=function(tree)
		return printnode(tree.root,'',true)
	end,
}

function splay.new(a)
	local root
	local Size
	
	if a then
		-- assumes SORTED array, if not sorted then call table.sort BEFORE passing it
		Size=#a
		root=(Size>0 and dc(a,1,Size))or nil
	else
		Size=0
	end
	
	local self={
		root=root,
		Size=Size,
		
		empty=empty,
		insert=insert,
		deletenode=deletenode,
		deletedata=deletedata,
		find=find,
		lowerbound=lowerbound,
		upperbound=upperbound,
	}
	
	setmetatable(self,splay)
	return self
end

return splay
2 Likes
#5

vEB-van Emde Boas Tree

I don’t really use this and its unlikely that I ever will so I am releasing it here
The code is not perfectly written because I wrote a little more than a year ago
Hopefully somebody finds it useful

--[[
	vEB (van Emde Boas Trees)
	
	This is slow because it recurses and calls functions. Because of this, it is recommended to use a Red-Black Binary Search Tree instead.
	
    u defines the desired upper bound, each node has u^0.5 children.
    You want u to be as close as possible to the maximum Value.

    You can insert negative numbers as well but it is a bit inefficient so it is better to offset all of the numbers by a constant beforehand.
	
	Sources:
		https://ocw.mit.edu/courses/electrical-engineering-and-computer-science/6-046j-design-and-analysis-of-algorithms-spring-2015/lecture-notes/MIT6_046JS15_lec04.pdf
		https://www.youtube.com/watch?v=hmReJCupbNU
		https://github.com/erikwaing/VEBTree
		
	Complexity:
		Min: O(1)
		Max: O(1)
		Member: O(lglg(u))
		Upper Bound: O(lglg(u))
		Lower Bound: O(lglg(u))
		Insertion: O(lglg(u))
		Deletion: O(lglg(u))

	Example Usage:
		local v = require(this).New(15)

		v:Insert(1)
		v:Insert(3)
		v:Insert(4)
		v:Insert(5)
		v:Insert(7)
		v:Insert(8)
		v:Insert(10)
		v:Insert(0)
		v:Insert(12)
		v:Insert(14)
		
		print(v:UpperBound(13))
		
		-- iterating
		local i = v.Min
		while i do
			print(i)
			i = v:Successor(i)
		end
	
--]]

-- Constants

-- efficiency variables
local error=error
local floor=math.floor

-- local variables

local function newveb(u)
	local self={
		u,-- u
		false,-- min
		false,-- max
		false,
		false,
		floor(u^.5),
	}
	
	if u>2 then
		local clusters={}
		for i=1,self[6]+1 do
			clusters[i]=false
		end
		
		self[4]=clusters-- clusters
		--self[5] = nil -- summary
	end
	
	return self
end

local function high(self,x)
	return floor(x/self[6])+1
end

local function low(self,x)
	return x%self[6]+1
end

local function index(self,h,l)
	return(h-1)*self[6]+l-1
end

local function member(self,x)
	-- returns whether x is a member of the vEB
	if x==self[2]or x==self[3]then
		-- found x as the min or max
		return true
	elseif self[1]<=2 then
		-- has not found x in the "leaf"
		return
	else
		local cluster=self[4][high(self,x)]
		if cluster then
			-- look for x in the cluster
			return member(cluster,low(self,x))
		end
	end
end

local function successor(self,x)
	local m=self[2]
	if self[1]<=2 then
		if x==0 and self[3]==1 then
			return 1
		end
	elseif m and x<m then
		-- x is less than everything in the tree, return the minimum
		return m
	else
		local h=high(self,x)
		local l=low(self,x)
		m=self[4]
		local cluster=m[h]
		local n=cluster and cluster[3]
		
		if n and l<n then
			return index(self,h,successor(cluster,l))
		else
			n=self[5]
			if n then
				n=successor(n,h)
				if n then
					local cluster2=m[n]
					if cluster2 then
						return index(self,n,cluster2[2])
					end
				end
			end
		end
	end
end

local function predecessor(self,x)
	local n=self[2]
	local m=self[3]
	if self[1]<=2 then
		if x==1 and n==0 then
			return 0
		end
	elseif m and x>m then
		return m
	else
		local h=high(self,x)
		local l=low(self,x)
		m=self[4]
		local cluster=m[h]
		local minlow=cluster and cluster[2]
		
		if minlow and l>minlow then
			return index(self,h,predecessor(cluster,l)or 0)
		else
			local predcluster=self[5]
			if predcluster then
				predcluster=predecessor(predcluster,h)
			end
			if predcluster then
				m=m[predcluster]
				return index(self,predcluster,(m and m[3])or 0)
			elseif n and x>n then
				return n
			end
		end
	end
end

local function insert(self,x)
	local m=self[2]
	if m then
		if x<m then
			-- swap x and min if x is less than min to *actually* insert min instead of x
			self[2]=x
			x=m
		end
		
		m=self[1]
		if m>2 then
			m=high(self,m)
			local h=high(self,x)
			local cluster=self[4][h]
			local summary=self[5]
			
			if not cluster then
				-- if the cluster doesn't exist then create it (save space by not having it exist) and set its bound to sqrt of u
				cluster=newveb(m)
				self[4][h]=cluster
			end
			if not summary then
				summary=newveb(m)
				self[5]=summary
			end
			if cluster[2]then
				insert(cluster,low(self,x))
			else
				insert(summary,h)
				local x=low(self,x)
				cluster[2]=x
				cluster[3]=x
			end
		end
		if x>self[3]then
			self[3]=x
		end
	else
		self[2]=x
		self[3]=x
	end
end

local function delete(self,x)
	local h
	local m=self[4]
	local p=self[5]
	
	if x==self[2]then
		h=p
		h=h and h[2]
		if not h then
			self[2]=nil
			self[3]=nil
			return
		end
		
		x=index(self,h,m[h][2])
		self[2]=x
	end
	
	local o=high(self,x)
	local n=m[o]
	delete(n,low(self,x))
	if not n[2]then
		delete(p,o)
		m[o]=false--nil -- delete it if its not being used
	end
	if x==self[3]then
		local h=p[3]
		self[3]=h and index(self,h,m[h][3])or self[2]
	end
end

local function vinsert(self,x)
	insert(self,x)
	self.min=self[2]
	self.max=self[3]
	self.Size=self.Size+1
	
	return self
end

local function vdelete(self,x)
	delete(self,x)
	self.min=self[2]
	self.max=self[3]
	self.Size=self.Size-1
	
	return self
end

local function vupperbound(self,x)
	return predecessor(self,x+1)
end

local function vlowerbound(self,x)
	return successor(self,x-1)
end

local function vclear(self)
	self[2]=false
	self[3]=false
	local a=self[4]
	if a then
		for i=1,#a do
			a[i]=false
		end
	end
	self[5]=false
	
	return self
end

local function vempty(self)
	return not self[1]
end

local veb={
	insert=vinsert,
	delete=vdelete,
	member=member,
	upperbound=vupperbound,
	lowerbound=vlowerbound,
	successor=successor,
	predecessor=predecessor,
	clear=vclear,
	empty=vempty,
}

function veb.new(u)
	if u<0 then
		error('Upper bound cannot be less than 0',2)
	end
	
	local self=newveb(u)
	
	self.Size=0
	self.max=false
	self.min=false
	self.u=u
	
	self.insert=vinsert
	self.delete=vdelete
	self.member=member
	self.upperbound=vupperbound
	self.lowerbound=vlowerbound
	self.clear=vclear
	self.empty=vempty
	self.successor=successor
	self.predecessor=predecessor
	
	return self
end

return veb
5 Likes