Chương trình tìm khoảng cách giữa hai nút trong cây nhị phân trong Python

Mô tả bài toán

Cho cây nhị phân và hai nút. Tìm khoảng cách (số cạnh) giữa hai nút này.

Ví dụ

Cây: [5, 3, 7, 2, 4, 6, 8] (level order). Nút 2 và 8 -> khoảng cách = 4 (đường đi: 2-3-5-7-8).

Giải thuật

1. Tìm LCA (Lowest Common Ancestor) của hai nút.
2. Tính khoảng cách từ LCA đến từng nút bằng BFS.
3. Tổng hai khoảng cách là kết quả.

Triển khai Python

import collections

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.data = data
        self.left = left
        self.right = right

def findLca(root, p, q):
    if root is None:
        return None
    if root.data in (p, q):
        return root
    left = findLca(root.left, p, q)
    right = findLca(root.right, p, q)
    if left and right:
        return root
    return left or right

def findDist(root, data):
    queue = collections.deque()
    queue.append((root, 0))
    while queue:
        current, dist = queue.popleft()
        if current.data == data:
            return dist
        if current.left:
            queue.append((current.left, dist + 1))
        if current.right:
            queue.append((current.right, dist + 1))

def solve(root, p, q):
    node = findLca(root, p, q)
    return findDist(node, p) + findDist(node, q)

# Helper functions to build tree omitted for brevity
root = make_tree([5, 3, 7, 2, 4, 6, 8])
print(solve(root, 2, 8))

Đầu vào

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

Đầu ra

4