Kiểm tra cây nhị phân có cân bằng chiều cao như cây Đỏ-Đen trong Python

Giả sử có một cây Đỏ-Đen (Red-Black Tree), ở đây chiều cao lớn nhất của một nút tối đa gấp đôi chiều cao nhỏ nhất. Nếu chúng ta có một cây tìm kiếm nhị phân, ta cần kiểm tra tính chất sau: Đối với mọi nút, độ dài đường đi dài nhất từ lá đến nút không quá gấp đôi số nút trên đường đi ngắn nhất từ nút đến lá.

Ví dụ, nếu đầu vào như hình dưới đây:

thì kết quả sẽ là True, vì cây này cân bằng.

Thuật toán giải quyết

Để giải quyết bài toán này, ta thực hiện các bước sau:

• Định nghĩa hàm solve() nhận tham số root, max_height, min_height.
• Nếu root là null:
  • Gán max_height := 0, min_height := 0.
  • Trả về True.
• Khởi tạo left_max := 0, left_min := 0, right_max := 0, right_min := 0.
• Gọi đệ quy solve(root.left, left_max, left_min). Nếu trả về False thì trả về False.
• Gọi đệ quy solve(root.right, right_max, right_min). Nếu trả về False thì trả về False.
• Cập nhật max_height := max(left_max, right_max) + 1.
• Cập nhật min_height := min(left_min, right_min) + 1.
• Nếu max_height <= 2 * min_height thì trả về True, ngược lại trả về False.
• Từ hàm chính, khởi tạo max_height := 0, min_height := 0 và gọi solve(root, max_height, min_height).

Ví dụ minh họa

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

def solve(root, max_height, min_height):
    if root is None:
        max_height = min_height = 0
        return True
    left_max = left_min = 0
    right_max = right_min = 0
    if solve(root.left, left_max, left_min) == False:
        return False
    if solve(root.right, right_max, right_min) == False:
        return False
    max_height = max(left_max, right_max) + 1
    min_height = min(left_min, right_min) + 1
    if max_height <= 2 * min_height:
        return True
    return False

def is_tree_balanced(root):
    max_height, min_height = 0, 0
    return solve(root, max_height, min_height)

root = TreeNode(10)
root.left = TreeNode(5)
root.right = TreeNode(100)
root.right.left = TreeNode(50)
root.right.right = TreeNode(150)
root.right.left.left = TreeNode(40)
print(is_tree_balanced(root))  # True

Đầu vào

root = TreeNode(10)
root.left = TreeNode(5)
root.right = TreeNode(100)
root.right.left = TreeNode(50)
root.right.right = TreeNode(150)
root.right.left.left = TreeNode(40)

Đầu ra

True