Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> Python

Tìm trung vị của BST trong O(n) thời gian và O(1) không gian trong Python

Cho cây nhị phân tìm kiếm (BST), tìm trung vị. Với số nút chẵn: median = (n/2-th + (n+1)/2-th) / 2. Với số nút lẻ: median = (n+1)/2-th node.

Ví dụ: cây BST với các giá trị 2,3,4,5,6,7,8,9,10 → Trung vị = 7.

Thuật toán Morris Inorder Traversal

Duyệt inorder không dùng stack/đệ quy (O(1) space):

  1. Đếm tổng số nút node_count bằng Morris traversal.
  2. Duyệt lại Morris, đếm count_curr:
  3. Nếu số nút lẻ và count_curr == (node_count+1)//2 → trả về giá trị hiện tại.
  4. Nếu số nút chẵn và count_curr == node_count//2 + 1 → trả về trung bình của previous và current.

Ví dụ code

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

def number_of_nodes(root):
    node_count = 0
    if root is None:
        return node_count
    current = root
    while current:
        if current.left is None:
            node_count += 1
            current = current.right
        else:
            previous = current.left
            while previous.right and previous.right != current:
                previous = previous.right
            if previous.right is None:
                previous.right = current
                current = current.left
            else:
                previous.right = None
                node_count += 1
                current = current.right
    return node_count

def calculate_median(root):
    if root is None:
        return 0
    node_count = number_of_nodes(root)
    count_curr = 0
    current = root
    previous = None
    while current:
        if current.left is None:
            count_curr += 1
            if node_count % 2 != 0 and count_curr == (node_count + 1) // 2:
                return previous.data
            elif node_count % 2 == 0 and count_curr == (node_count // 2) + 1:
                return (previous.data + current.data) // 2
            previous = current
            current = current.right
        else:
            previous = current.left
            while previous.right and previous.right != current:
                previous = previous.right
            if previous.right is None:
                previous.right = current
                current = current.left
            else:
                previous.right = None
                previous = previous
                count_curr += 1
                if node_count % 2 != 0 and count_curr == (node_count + 1) // 2:
                    return current.data
                elif node_count % 2 == 0 and count_curr == (node_count // 2) + 1:
                    return (previous.data + current.data) // 2
                previous = current
                current = current.right

# Tạo cây:       7
#            4     9
#           2 5   8 10
root = TreeNode(7)
root.left = TreeNode(4)
root.right = TreeNode(9)
root.left.left = TreeNode(2)
root.left.right = TreeNode(5)
root.right.left = TreeNode(8)
root.right.right = TreeNode(10)
print(calculate_median(root))

Input

BST với 7 nút: 2,3,4,5,6,7,8,9,10

Output

7