AVLtree/tree.cpp

#include<fstream>
#include<iostream>
#include<vector>
#include<memory>

template <typename T>
struct Node;

template <typename T>
using Node_ptr = std::unique_ptr<Node<T>>;

template <typename T>
struct Node{
    T value;
    int max_depth;
    Node_ptr<T> small;
    Node_ptr<T> big;

    bool add(T val);
    bool add_to_child(Node_ptr<T>* child, T val);
    void update_depth();
    void dump(std::ofstream& ofs) const;
    Node(T v): value{v}, max_depth{0}, small{nullptr}, big{nullptr}{}
};

template <typename T>
class Tree{
    Node_ptr<T> root;
    Node_ptr<T>* find(T val);
    bool try_to_del(Node_ptr<T>* node, T val);

public:
    bool insert(T val);
    bool del(T val);
    void print() const;
};

template <typename T>
Node_ptr<T>* Tree<T>::find(T val){
    for(Node_ptr<T>* n = &root; *n;){
        if(val == (*n)->value) return n;
        else if(val < (*n)->value) n = &(*n)->small;
        else n = &(*n)->big;
    }
    return nullptr;
}

template <typename T>
bool Tree<T>::insert(T val){
    if(!root){
        root = std::make_unique<Node<T>>(val);
        return true;
    }
    else return root->add(val);
}

template <typename T>
bool Node<T>::add(T val){
    if (val == value)     return false;
    else if (val < value) return add_to_child(&small, val);
    else                  return add_to_child(&big, val);
}

template <typename T>
bool Node<T>::add_to_child(Node_ptr<T>* child, T val){
    if(!(*child)){
        *child = std::make_unique<Node<T>>(val);
        update_depth();
        return true;
    }
    else{
        if((*child)->add(val)){
            update_depth();
            return true;
        }
        return false;
    }
}


template <typename T>
bool Tree<T>::del(T val){
    if(!root) return false;
    else      return try_to_del(&root, val);
}

template <typename T>
bool Tree<T>::try_to_del(Node_ptr<T>* node, T val){
    if(val == (*node)->value) {
        Node<T>* temp = (*node).get();
        std::vector<Node<T>*> path;
        if((*node)->small) {
            temp = (*node)->small.get();
            path.push_back(temp);
        }
        else {
            node->reset((*node)->big.release());
            return true;
        }
        while(temp->big){
            temp = temp->big.get();
            path.push_back(temp);
        }
        //mam temp na 'to swap' nodu...
        std::swap((*node)->value, temp->value);
        //temp->steal_values(std::move(temp->small));
        for(int i = path.size() -1; i >= 0; i--) path[i]->update_depth();
        path.clear();
        return true;
    }
    else if(val < (*node)->value){
        if(!(*node)->small) return false;
        else {
            if(try_to_del(&(*node)->small, val)){
                (*node)->update_depth();
                return true;
            }
            return false;
        }
    }
    else {
        if(!(*node)->big) return false;
        else {
            if(try_to_del(&(*node)->big, val)){
                (*node)->update_depth();
                return true;
            }
            return false;
        }
    }
}

template <typename T>
void Node<T>::update_depth(){
    if(small && big) max_depth = std::max(small->max_depth, big->max_depth) + 1;
    else if (small) max_depth = small->max_depth + 1;
    else max_depth = big->max_depth + 1;
}

template <typename T>
void Node<T>::dump(std::ofstream& ofs) const {
    if (small) {
        ofs << '"' << value << " : " << max_depth << '"' << " -> "  << '"' << small->value << " : " << small->max_depth << '"' << '\n';
        small->dump(ofs);
    }
    if (big) {
        ofs << '"' << value << " : " << max_depth << '"' << " -> "  << '"' << big->value << " : " << big->max_depth << '"' << '\n';
        big->dump(ofs);
    }
}

template <typename T>
void Tree<T>::print() const {
    std::ofstream ofs;
    ofs.open("viz.dot");
    ofs << "digraph Tree {\n";
    if(root) root->dump(ofs);
    ofs << "}";

    ofs.close();
    system("cat viz.dot | dot -Tx11");
}

int main(){
    Tree<int> t;
    std::vector<int> v = {3, 6, 8};
    for(auto i : v){
        if(t.insert(i)) std::cout << "added " << i << "\n";
        else std::cout << "failed at adding " << i << "\n";
        t.print();
    }
    for(auto i : v){
        if(t.del(i)) std::cout << "deleted " << i << "\n";
        else std::cout << "failed at deleting " << i << "\n";
        t.print();
    }
}