230. Kth Smallest Element in a BST Follow Up Solution (B+ Tree)

cplusplus-leetcode-230.cpp

/**
Hence without any optimisation insert/delete + search of kth element has O(2H + k) complexity. 
How do we optimise it ?

First, we need to make sure the balanced tree for sure. Here, we don't need to implement it (AVL, Red-Black Tree)
Secondly, we can mimic B+trees (https://en.wikipedia.org/wiki/B+ tree).
we can add a doubly linked list trait on the tree node. Then, as usual, 
we insert node with time complexity O(H), but take time complexity O(1) to maintain doubly linked list.
On the other hand, we can get the same complexity on deletion. 
**/
template <typename T>
struct BST {
    struct Node {
        T val;
        shared_ptr<Node> left, right, prev, next; // prev, next for doubly linked list
        Node() {}
        Node(T val) : val(val) {}

    };

    BST(T val) {
        root = make_shared<Node>(val);
        head = root;
        tail = root;
    }

    shared_ptr<Node> root, head, tail;

    void insert(T val);
    void insert_before(shared_ptr<Node> target, shared_ptr<Node> ptr); 
    void insert_after(shared_ptr<Node> target, shared_ptr<Node> ptr);
    void del(shared_ptr<Node> target);
    void del(T val);
    vector<T> inorder();
    vector<T> display();
    T kthSmallest(int k);
};



template<typename T>
vector<T> BST<T>::inorder() {
    vector<T> rst;
    auto cur = BST<T>::root;
    stack<shared_ptr<Node>> st;
    while(cur || !st.empty()) {
        if(cur) {
            st.push(cur);
            cur = cur->left;
        } else {
            rst.push_back(st.top()->val);
            cur = st.top()->right; st.pop();
        }
    }
    return rst;
}

template<typename T>
vector<T> BST<T>::display() {
    vector<T> rst;
    auto cur = head;
    while(cur) {
        rst.push_back(cur->val);
        cur = cur->next;
    }
    return rst;
}

template<typename T>
void BST<T>::insert(T val) {
   
    function<shared_ptr<Node>(shared_ptr<Node>, T)> helper = [&](shared_ptr<Node> ptr, T val) {
        if(!ptr) return make_shared<Node>(val);
        if(val < ptr->val) {
            bool isleaf = ptr->left == nullptr;
            ptr->left = helper(ptr->left, val);
            if(isleaf) insert_before(ptr, ptr->left);
        } else {
            bool isleaf = ptr->right == nullptr;
            ptr->right = helper(ptr->right, val);
            if(isleaf) insert_after(ptr, ptr->right);
        }
        return ptr;
    };
    helper(this->root, val);
}

template<typename T>
void BST<T>::del(T val) {
    function<shared_ptr<Node>(shared_ptr<Node>)> find_nxt = [](shared_ptr<Node> ptr) {
        while(ptr && ptr->left)
            ptr = ptr->left;
        return ptr;
    };
    function<shared_ptr<Node>(shared_ptr<Node>, T)> helper= [&](shared_ptr<Node> ptr, T val) {
        if(!ptr) return ptr;
        if(val < ptr->val) ptr->left = helper(ptr->left, val);
        else if ( val > ptr->val) ptr->right = helper(ptr->right, val);
        else {
            if(!ptr->left && !ptr->right) {
                ptr = nullptr;
                return ptr;
            }
            if(!ptr->left) {
                auto tmp = ptr->right;
                ptr = nullptr;
                return tmp;
            }
            if(!ptr->right) {
                auto tmp = ptr->left;
                ptr = nullptr;
                return tmp;
            }
            auto nxt = find_nxt(ptr->right);
            ptr->val = nxt->val;
            ptr->right = helper(ptr->right, ptr->val);
        }
        return ptr;
    };
    function<shared_ptr<Node>(shared_ptr<Node>, T)> helper2= [&](shared_ptr<Node> ptr, T val) {
        if(!ptr) return ptr;
        if(val < ptr->val) ptr->left = helper2(ptr->left, val);
        else if (val > ptr->val) ptr->right = helper2(ptr->right, val);
        else {
            del(ptr);
            if(!ptr->left && !ptr->right) {
                ptr = nullptr;
                return ptr;
            }
            if(!ptr->left) {
                auto tmp = ptr->right;
                ptr = nullptr;
                return tmp;
            }
            if(!ptr->right) {
                auto tmp = ptr->left;
                ptr = nullptr;
                return tmp;
            }
            auto nxt = find_nxt(ptr->right);
            ptr->val = nxt->val;
            ptr->right = helper2(ptr->right, ptr->val);
        }
        return ptr;
    };
    helper2(root, val);
}

template<typename T>
void BST<T>::insert_before(shared_ptr<BST::Node> target, shared_ptr<BST::Node> ptr) {
    if(target == head) {
        head->prev = ptr;
        ptr->next = head;
        head = ptr;
    } else {
        auto prev = target->prev;
        prev->next =  ptr;
        ptr->prev = prev;

        ptr->next = target;
        target->prev = ptr;
    }
}

template<typename T>
void BST<T>::insert_after(shared_ptr<BST::Node> target, shared_ptr<BST::Node> ptr) {
    if(target == tail) {
        tail->next = ptr;
        ptr->prev = tail;
        tail = ptr;
    } else {
        auto nxt = target->next;
        target->next = ptr;
        ptr->prev = target;
        nxt->prev = ptr;
        ptr->next = nxt;
    }
}

template<typename T>
void BST<T>::del(shared_ptr<BST::Node> target) {
    if ( head == tail ) {
        head = nullptr;
        tail = nullptr;
    } else if(tail == target) {
        tail = target->prev;
        tail->next = nullptr;
    } else if (head == target) {
        head = head->next;
        head->prev = nullptr;
    } else {
        auto prev = target->prev;
        auto nxt = target->next;
        prev->next = nxt;
        nxt->prev = prev;
    }
}
template<typename T>
T BST<T>::kthSmallest(int k) {
    auto cur = head;
    while(cur) {
        if(--k==0) return cur->val;
        cur = cur->next;
    }
}


TEST_CASE("BST insert delete") {

    BST<int> bst(5);
    bst.insert(3);
    bst.insert(2);
    bst.insert(4);
    bst.insert(7);
    bst.insert(6);
    bst.insert(8);

    vector<int> ans = {2,3,4,5,6,7,8};
    printf("Inorder traversal of the given tree \n");
    REQUIRE(ans  ==  bst.inorder());
    REQUIRE(ans  ==  bst.display());
    printf("\nDelete 2\n");
    bst.del(2);
    printf("Inorder traversal of the modified tree \n");
    ans = {3,4,5,6,7,8};
    REQUIRE(ans  ==  bst.inorder());
    REQUIRE(ans  ==  bst.display());
    REQUIRE( 5 == bst.kthSmallest(3));

    printf("\nDelete 3\n");
    bst.del(3);
    printf("Inorder traversal of the modified tree \n");
    ans = {4,5,6,7,8};
    REQUIRE(ans  ==  bst.inorder());
    REQUIRE(ans  ==  bst.display());

    printf("\nDelete 6\n");
    bst.del(6);
    printf("Inorder traversal of the modified tree \n");
    ans = {4,5,7,8};
    REQUIRE(ans  ==  bst.inorder());
    REQUIRE(ans  ==  bst.display());

    printf("\nInsert 6\n");
    bst.insert(6);
    printf("Inorder traversal of the modified tree \n");
    ans = {4,5,6,7,8};
    REQUIRE(ans  ==  bst.inorder());
    REQUIRE(ans  ==  bst.display());

    REQUIRE( 6 == bst.kthSmallest(3));
    REQUIRE( 5 == bst.kthSmallest(2));
};