准备刷LeetCode树专题的时候,发现要测试自己写的代码能不能给出正确结果的时候,必须要把题目给出的元素列表先构建成二叉树,这样自己在本地的时候才方便测试代码是不是对的。所以自己先把树的构建代码写出来,以后刷题直接拷贝来用比较方便。
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
class Tree:
def __init__(self):
self.root = None # 初始化时根节点为None
self.queue = [] # 用来保存子树未添加完的节点
def add(self, elem):
# 给树添加节点
node = TreeNode(elem) # 首先有一个node节点来保存每一要添加的节点
if self.root is None: # 如果树还未构建,即根节点root为空的话, 则构建根节点并且把根节点加入queue中
self.root = node
self.queue.append(self.root)
else:
treeNode = self.queue[0] # 根节点已经构建好的情况下,取出queue中第一个节点,为这个节点构建子树,注意这个时候是取出,不是弹出,弹出操作是在左右子树都构建完成时执行的
if treeNode.left is None:
treeNode.left = node
self.queue.append(treeNode.left) # 把这个节点的左子树节点加入queue中, 等待构建
else:
treeNode.right = node
self.queue.append(treeNode.right) # 把这个节点的右子树节点加入queue中, 等待构建
self.queue.pop(0) # 当treeNode的子树构建完毕时, 弹出这个已经构建好的节点,准备构建queue中下一个节点的子树(即queue[0])
if __name__ == '__main__':
tree = Tree()
nodeList = [1, 2, 3, 4, 5, 6, 7]
for node in nodeList:
tree.add(node) # 添加节点, 构建二叉树