Chương trình đếm số nút trong khoảng giá trị trên cây BST trong Python

Cho cây BST và hai giá trị l, r, đếm số nút có giá trị trong khoảng [l, r].

Ví dụ

Cây đầu vào: (hình ảnh), l=7, r=13.
Đầu ra: 3 (các nút 8, 10, 12).

Thuật toán

1. Sử dụng stack để duyệt cây (DFS).
2. Khởi tạo stack = [root], count = 0.
3. Khi stack không rỗng:
   • Lấy nút node.
   • Nếu node không null:
     • Nếu l <= node.data <= r: tăng count, đẩy cả hai con.
     • Nếu node.data < l: chỉ đẩy con phải.
     • Nếu node.data > r: chỉ đẩy con trái.
4. Trả về count.

Mã nguồn Python

from collections import deque

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

class Solution:
    def solve(self, root, l, r):
        stack, count = [root], 0
        while stack:
            node = stack.pop()
            if node:
                if l <= node.data <= r:
                    count += 1
                    stack += [node.right, node.left]
                elif node.data < l:
                    stack += [node.right]
                else:
                    stack += [node.left]
        return count

ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root, 7, 13))

Đầu vào / Đầu ra

Input: root = TreeNode(12) ..., 7,13
Output: 3