Chương trình tìm node bên phải trong cây nhị phân bằng Python

Cho cây nhị phân và node u. Tìm node ngay bên phải của u trên cùng level (u có thể là lá hoặc node nội bộ).

Thuật toán

• BFS level-order với deque
• Mỗi level: duyệt các node, lưu vào list temp, ghi index nếu gặp u
• Nếu u là node cuối level → return None
• Nếu tìm thấy u (index > -1): return temp[index+1]

Code Python

from queue import deque

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        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.val == element: return root
    res1 = search_node(root.left, element)
    if res1: return res1
    return search_node(root.right, element)

def solve(root, u):
    if not root: return None
    dq = deque([root])
    while dq:
        dq_size = len(dq)
        temp = []
        index = -1
        for _ in range(dq_size):
            node = dq.popleft()
            if node.left: dq.append(node.left)
            if node.right: dq.append(node.right)
            temp.append(node)
            if node == u:
                index = len(temp) - 1
        if index == len(temp) - 1:
            return None
        if index > -1:
            return temp[index + 1]
    return None

root = make_tree([5, 3, 7, 2, 4, 6, 8])
u = search_node(root, 6)
ret = solve(root, u)
print(ret.val)

Input

root = make_tree([5, 3, 7, 2, 4, 6, 8]); u = search_node(root, 6)

Output

8