Lua实现简单二叉树
定义二叉树
---@class BinaryTree 二叉树
---@field dep number 深度
---@field root BinaryTreeNode 根节点
---@field leafRootDic table 叶子节点字典
local BinaryTree = {}
- 二叉树的构造方法
---@return BinaryTree
function BinaryTree.New()
---@type BinaryTreeNode
local tree = {
dep = 0,
root = nil,
leafRootDic = {}
}
local meta = {
__index = BinaryTree
}
return setmetatable(tree, meta)
end
- 创建二叉树
---创建二叉树
function BinaryTree.CreateTree(dep, rootWeight, rootData)
dep = dep or 1
local tree = BinaryTree.New()
local root = BinaryTreeNode.New(
nil,
1,
1,
nil,
nil,
rootWeight,
rootData
)
root.isRootNode = true
tree.root = root
root.tree = tree
root:CreateChildNode(true, dep)
root:CreateChildNode(false, dep)
if not root.lChild and not root.rChild then
tree.leafRootDic[root:GetSeq()] = root
end
return tree
end
- 二叉树清理
---二叉树清理(该深度及以下的节点都会被销毁)
---@param dep number 清除的深度(该深度及以下的节点都会被销毁)
function BinaryTree:Clear(dep)
if dep == nil or dep < 2 then
dep = 2
end
self.root:Clear(dep)
end
- 销毁二叉树
---销毁二叉树
function BinaryTree:Destroy()
self.root:Clear(self.root:GetDeq())
self.dep = nil
self.root = nil
self.leafRootDic = nil
end
- 二叉树上的节点获取方法
---获取根节点
---@return BinaryTreeNode
function BinaryTree:GetRootNode()
return self.root
end
---@param node1 BinaryTreeNode
---@param node2 BinaryTreeNode
local function SortAllLeafNodeFunc(node1, node2)
return node1:GetSeq() < node2:GetSeq()
end
---获取所有的叶子节点
---@return BinaryTreeNode[]
function BinaryTree:GetAllLeafNode()
local allLeafList = {}
for i, leaf in pairs(self.leafRootDic) do
table.insert(allLeafList, leaf)
end
table.sort(allLeafList, SortAllLeafNodeFunc)
return allLeafList
end
- 打印二叉树信息
---打印二叉树信息
---@param func fun(data:any):string
function BinaryTree:Print(func)
self.root:Print(func)
end
---打印二叉树上所有的叶子信息
---@param func fun(data:any):string
function BinaryTree:PrintAllLeafNode(func)
local printStr = "AllLeafNode:\n"
for i, leafNode in ipairs(self:GetAllLeafNode()) do
printStr = printStr .. tostring(leafNode:GetSeq()) .. ":w=" .. tostring(leafNode:GetWeight()) .. (func and func(leafNode:GetData()) or "") .. " "
end
print(printStr .. "\n")
end
定义二叉树节点
---@class BinaryTreeNode 二叉树节点
---@field tree BinaryTree 二叉树
---@field seq number
---@field dep number 深度
---@field parent BinaryTreeNode 父节点
---@field isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@field lChild BinaryTreeNode 左子节点
---@field rChild BinaryTreeNode 右子节点
---@field weight number 权重
---@field data any 数据
---@field isRootNode boolean 是否为根节点
local BinaryTreeNode = {}
- 节点的构造方法
---@param tree BinaryTree 二叉树
---@param seq number
---@param dep number 深度
---@param parent BinaryTreeNode
---@param isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@param weight number 权重
---@param data any 数据
---@return BinaryTreeNode
function BinaryTreeNode.New(tree, seq, dep, parent, isLeft, weight, data)
---@type BinaryTreeNode
local node = {
tree = tree,
seq = seq or 1,
dep = dep or 1,
parent = parent,
isLeft = isLeft,
weight = weight,
data = data
}
local meta = {
__index = BinaryTreeNode
}
return setmetatable(node, meta)
end
- 创建子节点方法
---创建子节点
---@param isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@param dep number 深度
---@param weight number 权重
---@param data any 数据
---@return BinaryTreeNode
function BinaryTreeNode:CreateChildNode(isLeft, dep, weight, data)
dep = dep or self.dep + 1
if self.dep >= dep then
return nil
end
self:RemoveChildNode(isLeft)
---@type BinaryTreeNode
local new_node = BinaryTreeNode.New(
self.tree,
isLeft and self.seq * 2 or self.seq * 2 + 1,
self.dep + 1,
self,
isLeft,
weight,
data
)
if isLeft then
self.lChild = new_node
else
self.rChild = new_node
end
new_node:CreateChildNode(true, dep)
new_node:CreateChildNode(false, dep)
self.tree.leafRootDic[self.seq] = nil
if not new_node.lChild and not new_node.rChild then
self.tree.leafRootDic[new_node.seq] = new_node
end
return new_node
end
- 节点清理方法
---节点清理(该深度及以下的节点都会被销毁)
---@param dep number 清除的深度(该深度及以下的节点都会被销毁)
function BinaryTreeNode:Clear(dep)
dep = dep or self.dep + 1
if self.dep >= dep then
self:Destroy()
return
end
if self.lChild then
self.lChild:Clear(dep)
end
if self.rChild then
self.rChild:Clear(dep)
end
end
- 销毁节点方法
---销毁节点
function BinaryTreeNode:Destroy()
if self.lChild then
self.lChild:Destroy()
end
if self.rChild then
self.rChild:Destroy()
end
self.tree.leafRootDic[self:GetSeq()] = nil
if self.parent then
if self.isLeft then
self.parent.lChild = nil
else
self.parent.rChild = nil
end
if not self.parent.lChild and not self.parent.rChild then
self.tree.leafRootDic[self.parent:GetSeq()] = self.parent
end
end
self.tree = nil
self.seq = nil
self.dep = nil
self.parent = nil
self.isLeft = nil
self.lChild = nil
self.rChild = nil
self.weight = nil
self.data = nil
self.isRootNode = nil
end
- 节点所在的二叉树、节点的父/左子/右子节点的相关方法
---获取节点所在的二叉树
---@return BinaryTree
function BinaryTreeNode:GetTree()
return self.tree
end
---判断该节点是否存在父节点
---@return boolean
function BinaryTreeNode:CheckHasParentNode()
return self.parent ~= nil or not self.isRootNode
end
---获取父节点
---@return BinaryTreeNode
function BinaryTreeNode:GetParentNode()
return self.parent
end
---判断该节点是否为左节点（反之则为右节点）
---@return boolean
function BinaryTreeNode:GetIsLeft()
return self.isLeft
end
---判断该节点是否存在左/右子节点
---@return boolean
function BinaryTreeNode:CheckHasChildNode(isLeft)
return self:GetChildNode(isLeft) ~= nil
end
---获取左/右子节点
---@param isLeft boolean
---@return BinaryTreeNode
function BinaryTreeNode:GetChildNode(isLeft)
if isLeft then
return self.lChild
else
return self.rChild
end
end
---移除左/右子节点
function BinaryTreeNode:RemoveChildNode(isLeft)
if isLeft and self.lChild then
self.lChild:Destroy()
self.lChild = nil
elseif not isLeft and self.rChild then
self.rChild:Destroy()
self.rChild = nil
end
end
- 节点数据的相关方法
---@param weight number 权重
function BinaryTreeNode:SetWeight(weight)
self.weight = weight
end
---获取权重
---@return any
function BinaryTreeNode:GetWeight()
return self.weight
end
---计算路径权重 (当前节点往根节点方向的路径权重)
---@param func fun(data:any):number
---@return number
function BinaryTreeNode:GetPathWeight(func)
local pathWeight = self:CheckHasParentNode() and self:GetParentNode():GetPathWeight(func) or 0
local dataNum = func and func(self.data) or 1
return pathWeight + (self.weight or 0) * dataNum
end
---@param data any
function BinaryTreeNode:SetData(data)
self.data = data
end
---@return any
function BinaryTreeNode:GetData()
return self.data
end
---@return number
function BinaryTreeNode:GetSeq()
return self.seq
end
---@return number
function BinaryTreeNode:GetDeq()
return self.dep
end
---@param strList string[]
---@param func fun(data:any):string
function BinaryTreeNode:_GetPrintList(strList, func)
if not strList then
strList = {}
end
if strList[self.dep] then
strList[self.dep] = strList[self.dep] .. " " .. tostring(self.seq)
else
strList[self.dep] = tostring(self.seq)
end
strList[self.dep] = strList[self.dep] .. ":w=" .. tostring(self:GetWeight()) .. (func and func(self.data) or "") .. " "
if self.lChild then
self.lChild:_GetPrintList(strList, func)
end
if self.rChild then
self.rChild:_GetPrintList(strList, func)
end
return strList
end
- 打印节点信息方法
---打印节点信息
---@param func fun(data:any):string
function BinaryTreeNode:Print(func)
local printStr = "BinaryTree:\n"
local strList = self:_GetPrintList({}, func)
for i, str in ipairs(strList) do
printStr = printStr .. str .. "\n"
end
print(printStr)
end
完整代码
--region BinaryTree
---@class BinaryTree 二叉树
---@field dep number 深度
---@field root BinaryTreeNode 根节点
---@field leafRootDic table 叶子节点字典
BinaryTree = {}
---@return BinaryTree
function BinaryTree.New()
---@type BinaryTreeNode
local tree = {
dep = 0,
root = nil,
leafRootDic = {}
}
local meta = {
__index = BinaryTree
}
return setmetatable(tree, meta)
end
---创建二叉树
function BinaryTree.CreateTree(dep, rootWeight, rootData)
dep = dep or 1
local tree = BinaryTree.New()
local root = BinaryTreeNode.New(
nil,
1,
1,
nil,
nil,
rootWeight,
rootData
)
root.isRootNode = true
tree.root = root
root.tree = tree
root:CreateChildNode(true, dep)
root:CreateChildNode(false, dep)
if not root.lChild and not root.rChild then
tree.leafRootDic[root:GetSeq()] = root
end
return tree
end
---二叉树清理(该深度及以下的节点都会被销毁)
---@param dep number 清除的深度(该深度及以下的节点都会被销毁)
function BinaryTree:Clear(dep)
if dep == nil or dep < 2 then
dep = 2
end
self.root:Clear(dep)
end
---销毁二叉树
function BinaryTree:Destroy()
self.root:Clear(self.root:GetDeq())
self.dep = nil
self.root = nil
self.leafRootDic = nil
end
---获取根节点
---@return BinaryTreeNode
function BinaryTree:GetRootNode()
return self.root
end
---@param node1 BinaryTreeNode
---@param node2 BinaryTreeNode
local function SortAllLeafNodeFunc(node1, node2)
return node1:GetSeq() < node2:GetSeq()
end
---获取所有的叶子节点
---@return BinaryTreeNode[]
function BinaryTree:GetAllLeafNode()
local allLeafList = {}
for i, leaf in pairs(self.leafRootDic) do
table.insert(allLeafList, leaf)
end
table.sort(allLeafList, SortAllLeafNodeFunc)
return allLeafList
end
---@class BinaryTreeNode 二叉树节点
---@field tree BinaryTree 二叉树
---@field seq number
---@field dep number 深度
---@field parent BinaryTreeNode 父节点
---@field isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@field lChild BinaryTreeNode 左子节点
---@field rChild BinaryTreeNode 右子节点
---@field weight number 权重
---@field data any 数据
---@field isRootNode boolean 是否为根节点
BinaryTreeNode = {}
---@param tree BinaryTree 二叉树
---@param seq number
---@param dep number 深度
---@param parent BinaryTreeNode
---@param isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@param weight number 权重
---@param data any 数据
---@return BinaryTreeNode
function BinaryTreeNode.New(tree, seq, dep, parent, isLeft, weight, data)
---@type BinaryTreeNode
local node = {
tree = tree,
seq = seq or 1,
dep = dep or 1,
parent = parent,
isLeft = isLeft,
weight = weight,
data = data
}
local meta = {
__index = BinaryTreeNode
}
return setmetatable(node, meta)
end
---创建子节点
---@param isLeft boolean 是否为父节点的左子节点(反之则为右子节点)
---@param dep number 深度
---@param weight number 权重
---@param data any 数据
---@return BinaryTreeNode
function BinaryTreeNode:CreateChildNode(isLeft, dep, weight, data)
dep = dep or self.dep + 1
if self.dep >= dep then
return nil
end
self:RemoveChildNode(isLeft)
---@type BinaryTreeNode
local new_node = BinaryTreeNode.New(
self.tree,
isLeft and self.seq * 2 or self.seq * 2 + 1,
self.dep + 1,
self,
isLeft,
weight,
data
)
if isLeft then
self.lChild = new_node
else
self.rChild = new_node
end
new_node:CreateChildNode(true, dep)
new_node:CreateChildNode(false, dep)
self.tree.leafRootDic[self.seq] = nil
if not new_node.lChild and not new_node.rChild then
self.tree.leafRootDic[new_node.seq] = new_node
end
return new_node
end
---节点清理(该深度及以下的节点都会被销毁)
---@param dep number 清除的深度(该深度及以下的节点都会被销毁)
function BinaryTreeNode:Clear(dep)
dep = dep or self.dep + 1
if self.dep >= dep then
self:Destroy()
return
end
if self.lChild then
self.lChild:Clear(dep)
end
if self.rChild then
self.rChild:Clear(dep)
end
end
---销毁节点
function BinaryTreeNode:Destroy()
if self.lChild then
self.lChild:Destroy()
end
if self.rChild then
self.rChild:Destroy()
end
self.tree.leafRootDic[self:GetSeq()] = nil
if self.parent then
if self.isLeft then
self.parent.lChild = nil
else
self.parent.rChild = nil
end
if not self.parent.lChild and not self.parent.rChild then
self.tree.leafRootDic[self.parent:GetSeq()] = self.parent
end
end
self.tree = nil
self.seq = nil
self.dep = nil
self.parent = nil
self.isLeft = nil
self.lChild = nil
self.rChild = nil
self.weight = nil
self.data = nil
self.isRootNode = nil
end
--region 节点所在的二叉树、节点的父/左子/右子节点的相关方法
---获取节点所在的二叉树
---@return BinaryTree
function BinaryTreeNode:GetTree()
return self.tree
end
---判断该节点是否存在父节点
---@return boolean
function BinaryTreeNode:CheckHasParentNode()
return self.parent ~= nil or not self.isRootNode
end
---获取父节点
---@return BinaryTreeNode
function BinaryTreeNode:GetParentNode()
return self.parent
end
---判断该节点是否为左节点（反之则为右节点）
---@return boolean
function BinaryTreeNode:GetIsLeft()
return self.isLeft
end
---判断该节点是否存在左/右子节点
---@return boolean
function BinaryTreeNode:CheckHasChildNode(isLeft)
return self:GetChildNode(isLeft) ~= nil
end
---获取左/右子节点
---@param isLeft boolean
---@return BinaryTreeNode
function BinaryTreeNode:GetChildNode(isLeft)
if isLeft then
return self.lChild
else
return self.rChild
end
end
---移除左/右子节点
function BinaryTreeNode:RemoveChildNode(isLeft)
if isLeft and self.lChild then
self.lChild:Destroy()
self.lChild = nil
elseif not isLeft and self.rChild then
self.rChild:Destroy()
self.rChild = nil
end
end
--endregion
--region 节点数据的相关方法
---@param weight number 权重
function BinaryTreeNode:SetWeight(weight)
self.weight = weight
end
---获取权重
---@return any
function BinaryTreeNode:GetWeight()
return self.weight
end
---计算路径权重 (当前节点往根节点方向的路径权重)
---@param func fun(data:any):number
---@return number
function BinaryTreeNode:GetPathWeight(func)
local pathWeight = self:CheckHasParentNode() and self:GetParentNode():GetPathWeight(func) or 0
local dataNum = func and func(self.data) or 1
return pathWeight + (self.weight or 0) * dataNum
end
---@param data any
function BinaryTreeNode:SetData(data)
self.data = data
end
---@return any
function BinaryTreeNode:GetData()
return self.data
end
---@return number
function BinaryTreeNode:GetSeq()
return self.seq
end
---@return number
function BinaryTreeNode:GetDeq()
return self.dep
end
---@param strList string[]
---@param func fun(data:any):string
function BinaryTreeNode:_GetPrintList(strList, func)
if not strList then
strList = {}
end
if strList[self.dep] then
strList[self.dep] = strList[self.dep] .. " " .. tostring(self.seq)
else
strList[self.dep] = tostring(self.seq)
end
strList[self.dep] = strList[self.dep] .. ":w=" .. tostring(self:GetWeight()) .. (func and func(self.data) or "") .. " "
if self.lChild then
self.lChild:_GetPrintList(strList, func)
end
if self.rChild then
self.rChild:_GetPrintList(strList, func)
end
return strList
end
--endregion
--region Print
---打印节点信息
---@param func fun(data:any):string
function BinaryTreeNode:Print(func)
local printStr = "BinaryTree:\n"
local strList = self:_GetPrintList({}, func)
for i, str in ipairs(strList) do
printStr = printStr .. str .. "\n"
end
print(printStr)
end
---打印二叉树信息
---@param func fun(data:any):string
function BinaryTree:Print(func)
self.root:Print(func)
end
---打印二叉树上所有的叶子信息
---@param func fun(data:any):string
function BinaryTree:PrintAllLeafNode(func)
local printStr = "AllLeafNode:\n"
for i, leafNode in ipairs(self:GetAllLeafNode()) do
printStr = printStr .. tostring(leafNode:GetSeq()) .. ":w=" .. tostring(leafNode:GetWeight()) .. (func and func(leafNode:GetData()) or "") .. " "
end
print(printStr .. "\n")
end
--endregion
--endregion
测试输出：
local function TestBinaryTree()
local binaryTree = BinaryTree.CreateTree(5)
binaryTree:Print()
binaryTree:PrintAllLeafNode()
--binaryTree:Clear(3)
--binaryTree:Print()
binaryTree:GetRootNode():GetChildNode(true):GetChildNode(true):Destroy()
binaryTree:Print()
binaryTree:PrintAllLeafNode()
binaryTree:Clear()
binaryTree:Print()
binaryTree:PrintAllLeafNode()
end
运行，得到：
BinaryTree:
1:w=nil
2:w=nil 3:w=nil
4:w=nil 5:w=nil 6:w=nil 7:w=nil
8:w=nil 9:w=nil 10:w=nil 11:w=nil 12:w=nil 13:w=nil 14:w=nil 15:w=nil
16:w=nil 17:w=nil 18:w=nil 19:w=nil 20:w=nil 21:w=nil 22:w=nil 23:w=nil 24:w=nil 25:w=nil 26:w=nil 27:w=nil 28:w=nil 29:w=nil 30:w=nil 31:w=nil
AllLeafNode:
16:w=nil 17:w=nil 18:w=nil 19:w=nil 20:w=nil 21:w=nil 22:w=nil 23:w=nil 24:w=nil 25:w=nil 26:w=nil 27:w=nil 28:w=nil 29:w=nil 30:w=nil 31:w=nil
BinaryTree:
1:w=nil
2:w=nil 3:w=nil
5:w=nil 6:w=nil 7:w=nil
10:w=nil 11:w=nil 12:w=nil 13:w=nil 14:w=nil 15:w=nil
20:w=nil 21:w=nil 22:w=nil 23:w=nil 24:w=nil 25:w=nil 26:w=nil 27:w=nil 28:w=nil 29:w=nil 30:w=nil 31:w=nil
AllLeafNode:
20:w=nil 21:w=nil 22:w=nil 23:w=nil 24:w=nil 25:w=nil 26:w=nil 27:w=nil 28:w=nil 29:w=nil 30:w=nil 31:w=nil
BinaryTree:
1:w=nil
AllLeafNode:
1:w=nil
验证成功，二叉树构造完成