Tìm cặp có tổng cho trước trong cây BST cân bằng bằng C++

Giả sử chúng ta có một cây tìm kiếm nhị phân (BST) cân bằng và một tổng mục tiêu, chúng ta cần xác định xem có tồn tại một cặp các nút trong cây có tổng bằng giá trị mục tiêu hay không. Lưu ý rằng cây BST là bất biến (immutable).

Ví dụ minh họa

Đối với cây BST như hình dưới đây và tổng mục tiêu là 35:

Kết quả sẽ là cặp (9 + 26 = 35).

Thuật toán

Chúng ta sử dụng hai ngăn xếp để duyệt cây theo thứ tự inorder (tăng dần) và reverse inorder (giảm dần) cùng lúc:

• Khởi tạo hai ngăn xếp s1 (cho inorder) và s2 (cho reverse inorder).
• Sử dụng hai con trỏ curr1 và curr2 bắt đầu từ gốc.
• Lặp vô hạn:
  • Đẩy các nút trái vào s1 cho đến khi gặp NULL, sau đó lấy nút trên cùng, gán giá trị cho val1, chuyển sang nhánh phải.
  • Đẩy các nút phải vào s2 cho đến khi gặp NULL, sau đó lấy nút trên cùng, gán giá trị cho val2, chuyển sang nhánh trái.
  • Nếu val1 != val2 và val1 + val2 == target, in cặp và trả về true.
  • Nếu tổng nhỏ hơn target, tiến con trỏ inorder (done1 = false).
  • Nếu tổng lớn hơn target, tiến con trỏ reverse inorder (done2 = false).
  • Nếu val1 >= val2, dừng lại và trả về false (đã duyệt hết cặp hợp lệ).

Triển khai C++

#include <bits/stdc++.h>
using namespace std;

class TreeNode {
public:
    int val;
    TreeNode *left, *right;
    TreeNode(int data) {
        val = data;
        left = NULL;
        right = NULL;
    }
};

bool isPairPresent(TreeNode* root, int target) {
    stack<TreeNode*> s1, s2;
    bool done1 = false, done2 = false;
    int val1 = 0, val2 = 0;
    TreeNode *curr1 = root, *curr2 = root;

    while (true) {
        // Inorder traversal (smallest to largest)
        while (!done1) {
            if (curr1 != NULL) {
                s1.push(curr1);
                curr1 = curr1->left;
            } else {
                if (s1.empty()) done1 = true;
                else {
                    curr1 = s1.top(); s1.pop();
                    val1 = curr1->val;
                    curr1 = curr1->right;
                    done1 = true;
                }
            }
        }

        // Reverse inorder traversal (largest to smallest)
        while (!done2) {
            if (curr2 != NULL) {
                s2.push(curr2);
                curr2 = curr2->right;
            } else {
                if (s2.empty()) done2 = true;
                else {
                    curr2 = s2.top(); s2.pop();
                    val2 = curr2->val;
                    curr2 = curr2->left;
                    done2 = true;
                }
            }
        }

        if (val1 != val2 && val1 + val2 == target) {
            cout << "Pair Found: " << val1 << " + " << val2 << " = " << target << endl;
            return true;
        } else if (val1 + val2 < target) {
            done1 = false;
        } else if (val1 + val2 > target) {
            done2 = false;
        }

        if (val1 >= val2) return false;
    }
}

int main() {
    TreeNode* root = new TreeNode(16);
    root->left = new TreeNode(11);
    root->right = new TreeNode(21);
    root->left->left = new TreeNode(9);
    root->left->right = new TreeNode(13);
    root->right->left = new TreeNode(17);
    root->right->right = new TreeNode(26);
    int target = 35;
    cout << (isPairPresent(root, target));
    return 0;
}

Kết quả

Pair Found: 9 + 26 = 35
1

Độ phức tạp: O(n) thời gian, O(h) không gian (h là chiều cao cây).