Cho cây nhị phân có lỗi: con trỏ right của một node trỏ sai sang node cùng level. Cần tìm node lỗi, xóa node đó và subtree của nó (trừ node bị trỏ sai). Trả về root đã sửa.
Thuật toán
- BFS với queue, map parent (p), set visited
- Khi gặp node đã visited: tìm grandparent và biết node lỗi là con trái/phải của grandparent → gán None
- Trả về root
Code Python
import collections
class TreeNode:
def __init__(self, data, left=None, right=None):
self.data = data
self.left = left
self.right = right
def insert(temp, data):
que = [temp]
while que:
temp = que.pop(0)
if not temp.left:
temp.left = TreeNode(data) if data is not None else TreeNode(0)
break
else:
que.append(temp.left)
if not temp.right:
temp.right = TreeNode(data) if data is not None else TreeNode(0)
break
else:
que.append(temp.right)
def make_tree(elements):
Tree = TreeNode(elements[0])
for element in elements[1:]:
insert(Tree, element)
return Tree
def search_node(root, element):
if not root: return None
if root.data == element: return root
res1 = search_node(root.left, element)
if res1: return res1
return search_node(root.right, element)
def print_tree(root):
if root:
print_tree(root.left)
print(root.data, end=', ')
print_tree(root.right)
def solve(root):
q = collections.deque([root])
p, visited = {}, set()
while q:
cur = q.popleft()
if cur in visited:
grand_p, is_left = p[p[cur][0]]
if is_left: grand_p.left = None
else: grand_p.right = None
return root
visited.add(cur)
if cur.left:
p[cur.left] = (cur, 1)
q.append(cur.left)
if cur.right:
p[cur.right] = (cur, 0)
q.append(cur.right)
return root
root = make_tree([5, 3, 7, 2, 4, 6, 8])
link_from = search_node(root, 4)
link_to = search_node(root, 6)
link_from.right = link_to
print_tree(solve(root))
Input
root = make_tree([5, 3, 7, 2, 4, 6, 8]); link_from.right = link_to
Output
2, 3, 5, 6, 7, 8,