写在前面
今天刷lc的每日一题, 碰到了关于二叉树的题目, 本来是简单层序遍历(BFS)出结果的, 但是我在本地测试的时候总是出问题, 层序遍历出来的结果竟然跟测试样例的二叉树长得不一样了!
看下面这个例子:
1
\
2
\
3
/ \
4 5
上面这颗树的层序遍历应该是:[1,None,2,None,3,4,5], 但是我的测试代码却得到了:
第1层元素: [1]
第2层元素: [None, 2]
第3层元素: [None, 3, 4, 5]
就是说直接忽略了空节点, 这可是大大影响了树的结构了.. 下面尝试解决这个问题.
代码(Python)
# 首先定义树的根节点
class Node(object):
"""docstring for Node"""
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
# 定义一棵二叉树
class BinaryTree(object):
"""docstring for BinaryTree"""
def __init__(self):
# 根节点
self.root = None
def add(self, item):
node = Node(item)
# 向树中插入元素, 使用队列存储元素, 读取与弹出
if not self.root:
self.root = node
return
# 用顺序表实现队列, 先入先出FIFO
queue = [self.root]
while queue:
cur_node = queue.pop(0)
# 若当前节点的左孩子为空, 将节点赋给当前节点左孩子
if not cur_node.left:
cur_node.left = node
return
else:
queue.append(cur_node.left)
if not cur_node.right:
cur_node.right = node
return
else:
queue.append(cur_node.right)
def breadth_travel1(self):
"""广度遍历: 方法同add, 是一种反过来的操作
"""
# 使用队列
q = [self.root]
if self.root is None:
return
level = 1
while q:
print(f"第{level}层元素: {[i.val for i in q]}")
nq = []
for cur_node in q:
if cur_node.left:
nq.append(cur_node.left)
if cur_node.right:
nq.append(cur_node.right)
q = nq
level += 1
if __name__ == '__main__':
tree = BinaryTree()
for i in [1, None, 2, None, 3, 4, 5]:
tree.add(i)
print("广度遍历: ")
tree.breadth_travel()
要解决的问题就是在构建树时(插入节点时)出现了None这个节点的时候, 就不要继续给树添加节点了, 而是直接往下遍历去找值不为None的结点进行插入, 有了这个思路, 我们可以把代码中的add函数加上一个判断, 如下:
def add(self, item):
node = Node(item)
# 向树中插入元素, 使用队列存储元素, 读取与弹出
if not self.root:
self.root = node
return
# 用顺序表实现队列, 先入先出FIFO
queue = [self.root]
while queue:
cur_node = queue.pop(0)
if cur_node.val is None: # 这里加一个判断
continue
# 若当前节点的左孩子为空, 将节点赋给当前节点左孩子
if not cur_node.left:
cur_node.left = node
return
else:
queue.append(cur_node.left)
if not cur_node.right:
cur_node.right = node
return
else:
queue.append(cur_node.right)
这时候再运行代码, 就可以得到:
第1层元素: [1]
第2层元素: [None, 2]
第3层元素: [None, 3]
第4层元素: [4, 5]