tage: Python BinaryTree DSA

寫在前面

今天刷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]