Не сбалансированное бинарное дерево поиска

unbalanced-search-binary-tree.cpp

#include <iostream>
#include <iterator>
#include <string>
#include <vector>

class Tree {
    struct TreeNode {
        TreeNode *left, *right;
        TreeNode *parent;

        string name;

        TreeNode(const string &name, TreeNode *parent = nullptr) : name(name), parent(parent) {
            left = right = nullptr;
        }

        int getHeight() const {
            int lh = left ? left->getHeight() : 0;
            int rh = right ? right->getHeight() : 0;
            return max(lh, rh) + 1;
        }
    };

    TreeNode *root;
    size_t size;

public:
    class iterator : public std::iterator<bidirectional_iterator_tag, string, string, const string*, string> {
        TreeNode *node;
    public:
        explicit iterator(TreeNode *node = nullptr) : node(node) {}
        iterator& operator++() {
            if (node->right) {
                node = node->right;
                while (node->left)
                    node = node->left;
            } else {
                while (node->parent && node == node->parent->right)
                    node = node->parent;
                node = node->parent;
            }
            return *this;
        }
        reference operator*() const { return node->name; }
        iterator operator++(int) {
            iterator retval = *this; ++(*this); return retval;
        }
        // TODO
    };

    Tree() : root(nullptr), size(0) {}

    size_t getSize() const { return size; }
    int getHeight() const { return root ? root->getHeight() : 0; }

    bool add(string name) {
        if (!root) {
            root = new TreeNode(name);
            size = 1;
            return true;
        }
        TreeNode *node = root;
        for (;;) {
            if (node->name == name)
                return false;
            if (node->name > name) {
                if (node->left == nullptr) {
                    TreeNode *newNode = new TreeNode(name, node);
                    node->left = newNode;
                    size++;
                    return true;
                }
                node = node->left;
            } else {
                if (node->right == nullptr) {
                    TreeNode *newNode = new TreeNode(name, node);
                    node->right = newNode;
                    size++;
                    return true;
                }
                node = node->right;
            }
        }
    }

    bool remove(string name) {
        TreeNode *node = root;
        while (node) {
            if (node->name == name) {
                TreeNode *replacement;
                if (node->left && node->right) {
                    replacement = node->right;
                    while (replacement->left)
                        replacement = replacement->left;
                    if (replacement->parent != node)
                        replacement->parent->left = replacement->right;
                    if (replacement->right)
                        replacement->right->parent = replacement->parent;
                    replacement->parent = node->parent;
                    replacement->right = node->right;
                    replacement->left = node->left;
                    replacement->left->parent = replacement;
                } else if (node->left) {
                    replacement = node->left;
                    node->left->parent = node->parent;
                } else if (node->right) {
                    replacement = node->right;
                    node->right->parent = node->parent;
                } else {
                    replacement = nullptr;
                }
                if (node->parent) {
                    if (node->parent->left == node)
                        node->parent->left = replacement;
                    else
                        node->parent->right = replacement;
                }
                if (node == root)
                    root = replacement;
                size--;
                delete node;
                return true;
            }
            if (node->name > name)
                node = node->left;
            else
                node = node->right;
        }
        return false;
    }

    bool contains(string name) const {
        TreeNode *node = root;
        while (node) {
            if (node->name == name)
                return true;
            if (node->name > name)
                node = node->left;
            else
                node = node->right;
        }
        return false;
    }
};

int main()
{
    Tree tree;
    vector<string> names = { "foo", "bar", "bar", "buz" };
    for (auto name : names)
        cout << "adding '" << name << "': " << tree.add(name) << "'" << endl;
    for (auto name : names)
        cout << "contains '" << name << "': " << tree.contains(name) << "'" << endl;
    for (auto name : names)
        cout << "remove '" << name << "': " << tree.remove(name) << "'" << endl;
}