Kiểm tra xem có bộ ba trong BST cân bằng cộng thành 0 hay không trong C++

Cho một cây tìm kiếm nhị phân cân bằng, viết hàm is_valid_triplet() trả về true nếu tồn tại bộ ba trong BST có tổng bằng 0, ngược lại false. Ràng buộc:

• Độ phức tạp thời gian O(n^2)
• Không gian thêm O(log n)

Ví dụ: với cây như hình, kết quả True vì bộ ba [-15, 7, 8].

Các bước:

• Định nghĩa bst_to_doubli_list(root, head, tail): chuyển BST thành danh sách liên kết đôi theo thứ tự in-order.
• Định nghĩa is_in_double_list(head, tail, sum): dùng hai con trỏ head/tail để tìm cặp có tổng = sum.
• Từ main: nếu root NULL trả về false; head=tail=NULL; gọi bst_to_doubli_list; duyệt head cho đến head->right != tail và head->key < 0: nếu is_in_double_list(head->right, tail, -head->key) trả về true; else head = head->right; cuối cùng trả về false.

Ví dụ (C++)

#include <bits/stdc++.h>
using namespace std;
class TreeNode {
public:
    int key;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : key(0), left(NULL), right(NULL) {}
    TreeNode(int x) : key(x), left(NULL), right(NULL) {}
};
void bst_to_doubli_list(TreeNode* root, TreeNode** head, TreeNode** tail) {
    if (root == NULL) return;
    if (root->left) bst_to_doubli_list(root->left, head, tail);
    root->left = *tail;
    if (*tail) (*tail)->right = root;
    else *head = root;
    *tail = root;
    if (root->right) bst_to_doubli_list(root->right, head, tail);
}
bool is_in_double_list(TreeNode* head, TreeNode* tail, int sum) {
    while (head != tail) {
        int current = head->key + tail->key;
        if (current == sum) return true;
        else if (current > sum) tail = tail->left;
        else head = head->right;
    }
    return false;
}
bool is_valid_triplet(TreeNode *root) {
    if (root == NULL) return false;
    TreeNode* head = NULL;
    TreeNode* tail = NULL;
    bst_to_doubli_list(root, &head, &tail);
    while ((head->right != tail) && (head->key < 0)){
        if (is_in_double_list(head->right, tail, -1*head->key)) return true;
        else head = head->right;
    }
    return false;
}
TreeNode* insert(TreeNode* root, int key) {
    if (root == NULL) return new TreeNode(key);
    if (root->key > key) root->left = insert(root->left, key);
    else root->right = insert(root->right, key);
    return root;
}
int main(){
    TreeNode* root = NULL;
    root = insert(root, 7);
    root = insert(root, -15);
    root = insert(root, 15);
    root = insert(root, -7);
    root = insert(root, 14);
    root = insert(root, 16);
    root = insert(root, 8);
    cout << is_valid_triplet(root);
}

Đầu vào

root = insert(root, 7); root = insert(root, -15); root = insert(root, 15); root = insert(root, -7); root = insert(root, 14); root = insert(root, 16); root = insert(root, 8);

Đầu ra

1