//
// Red-Black (balanced) tree
//
#ifndef RBTREE_H
#define RBTREE_H

// A node in Red-Black tree
class RBTreeNode {
public:
    RBTreeNode* left;   // pointer to the left son
    RBTreeNode* right;  // pointer to the right son
    RBTreeNode* parent; // pointer to the parent
    bool        red;    // the node is red (true) or black (false)
    void*       value;  // The value of tree node: normally, it is
                        //     a pair (key, value of key)
    RBTreeNode():
        left(0),
        right(0),
        parent(0),
        red(false),
        value(0)
    {
    }
};

typedef RBTreeNode* RBTreeNodePtr; // Pointer to the RBTreeNode

// An ABSTRACT class representing values in tree nodes
class RBTreeNodeValue {
public:
    RBTreeNodeValue() {}
    virtual ~RBTreeNodeValue() {}   // virtual destructor

    // Abstract class has virtual methods unimplemented

    // The values can be compared
    virtual int compareTo(const RBTreeNodeValue&) const = 0;

    bool operator==(const RBTreeNodeValue& k) const {
        return (compareTo(k) == 0);
    }
    bool operator!=(const RBTreeNodeValue& k) const {
        return (compareTo(k) != 0);
    }
    bool operator<(const RBTreeNodeValue& k) const {
        return (compareTo(k) < 0);
    }
    bool operator<=(const RBTreeNodeValue& k) const {
        return (compareTo(k) <= 0);
    }
    bool operator>(const RBTreeNodeValue& k) const {
        return (compareTo(k) > 0);
    }
    bool operator>=(const RBTreeNodeValue& k) const {
        return (compareTo(k) >= 0);
    }
};

// Every node of Red-Black tree is of red of black color.
// A leaf is an external NULL-node and has the black color.
// The root is black.
// The sons of red node must be black.
// Every path from the root to a leaf has the same number
// of black nodes (not including a root, but including a leaf).
class RBTree {
public:
    // Header contains a pointer to the root of the tree
    // as a left son.
    // The tree may be empty, in this case header->left == 0
    RBTreeNode  header;
    int numNodes;

    RBTree():
        header(),
        numNodes(0)
    {
        header.red = true;      // The header has the red color!
    }

    void clear();
    void erase() { clear(); }
    void removeAll() { clear(); }

    ~RBTree() { clear(); }

    RBTreeNode* root() { return header.left; }
    const RBTreeNode* root() const { return header.left; }

    int size() const { return numNodes; }

    // Find a key in a subtree
    // In: key         -- a key to find;
    //     subTreeRoot -- a root of a subtree. If subTreeRoot == 0,
    //                    then find in complete tree
    // Out: node       -- a pointer to the node that contains a key,
    //                    if key is in the set,
    //                    or a pointer to a node that should be parent to
    //                    the node with the key given.
    // Return value: true, if the subtree contains the key.
    bool find(
        const RBTreeNodeValue* key,
        const RBTreeNode* subTreeRoot = 0,
        RBTreeNodePtr* node = 0
    ) const;

    // Insert a key into the tree:
    //     create a new node and insert it as a leaf.
    // The color of a new node is red.
    // Should be called after the "find" method, that returned "false".
    void insert(
        RBTreeNode* parentNode,
        RBTreeNodeValue* v
    );

    // Rotate a node x to the left
    void rotateLeft(RBTreeNode* x);

    // Rotate a node x to the right
    void rotateRight(RBTreeNode* x);

    void rebalanceAfterInsert(RBTreeNode* x);

    // Remove a subtree and return the number of nodes removed
    int removeSubtree(RBTreeNode* subTreeRoot);

    const RBTreeNode* minimalNode(const RBTreeNode* subTreeRoot = 0) const;
    RBTreeNode* minimalNode(const RBTreeNode* subTreeRoot = 0) {
        return (RBTreeNode*)(
            ((const RBTree*) this)->minimalNode(subTreeRoot)
        );
    }

    const RBTreeNode* maximalNode(const RBTreeNode* subTreeRoot = 0) const;
    RBTreeNode* maximalNode(const RBTreeNode* subTreeRoot = 0) {
        return (RBTreeNode*)(
            ((const RBTree*) this)->maximalNode(subTreeRoot)
        );
    }

    const RBTreeNode* nextNode(const RBTreeNode* node) const;
    RBTreeNode* nextNode(RBTreeNode* node) {
        return (RBTreeNode*)(
            ((const RBTree*) this)->nextNode(node)
        );
    }

    const RBTreeNode* previousNode(const RBTreeNode* node) const;
    RBTreeNode* previousNode(RBTreeNode* node) {
        return (RBTreeNode*)(
            ((const RBTree*) this)->previousNode(node)
        );
    }

protected:
    void eraseNode(RBTreeNode* node);

public:
    class const_iterator {
    protected:
        const RBTree* tree;
        const RBTreeNode* node;
    public:
        const_iterator(): tree(0), node(0) {}

        const_iterator(
            const RBTree* t,
            const RBTreeNode* n
        ):
            tree(t),
            node(n)
        {}

        bool operator==(const const_iterator& i) const {
            return (tree == i.tree && node == i.node);
        }

        bool operator!=(const const_iterator& i) const {
            return !operator==(i);
        }

        const_iterator& operator++() {
            node = tree->nextNode(node);
            return *this;
        }

        const_iterator& operator--() {
            node = tree->previousNode(node);
            return *this;
        }

        const_iterator operator++(int) { // Post-increment (don't use it!)
            const_iterator tmp = *this;
            ++(*this);
            return tmp;
        }

        const_iterator operator--(int) { // Post-decrement (don't use it!)
            const_iterator tmp = *this;
            --(*this);
            return tmp;
        }

        const RBTreeNode& operator*() const { // Dereference
            return *node;
        }
        const RBTreeNode* operator->() const {
            return node;
        }
    };

    class iterator: public const_iterator {
    public:
        iterator():  
            const_iterator()
        {}

        iterator(RBTree* t, RBTreeNode* n):
            const_iterator(t, n)
        {}

        RBTreeNode& operator*() const { // Dereference
            return (RBTreeNode&)(
                ((const_iterator*) this)->operator*()
            );
        }
        RBTreeNode* operator->() const {
            return (RBTreeNode*)(
                ((const_iterator*) this)->operator->()
            );
        }
    };

    const_iterator begin() const {
        return const_iterator(this, minimalNode());
    }
    iterator begin() {
        return iterator(this, minimalNode());
    }

    const_iterator end() const {
        return const_iterator(this, &header);
    }
    iterator end() {
        return iterator(this, &header);
    }
};

#endif /* RBTREE_H */

// Red-Black Tree
// class RBTree, implementation
#include <assert.h>
#include "RBTree.h"

// Find a key in a subtree
// In: key  -- a key to find;
//     root -- a root of a subtree. If root == 0,
//             then find in complete tree
// Out: node -- a pointer to the node that contains a key,
//                  if key is in the set,
//              or a pointer to a node that should be parent to
//                  a node with this key in case of insertion
// Return value: true, if the subtree contains the key.
bool RBTree::find(
    const RBTreeNodeValue* key,
    const RBTreeNode* subTreeRoot /* = 0 */,
    RBTreeNodePtr* node /* = 0 */
) const {
    const RBTreeNode* x;        // current node
    const RBTreeNode* y;        // its parent
    if (subTreeRoot == 0) {
        x = root();
        y = &header;
    } else {
        x = subTreeRoot;
        y = x->parent;
    }
    while (x != 0) {
        const RBTreeNodeValue* currentKey = (const RBTreeNodeValue*) x->value;
        int n = key->compareTo(*currentKey);
        y = x;
        if (n == 0) {
            // key is found
            if (node != 0)
                *node = (RBTreeNode*) x;
            return true;
        } else if (n < 0) {
            x = x->left;
        } else {
            x = x->right;
        }
    }

    // key is not in the tree
    if (node != 0)
        *node = (RBTreeNode*) y;
    return false;
}

// Insert a key into the tree:
//     create a new node and insert it as a leaf.
// The color of a new node is red.
// Should be called after the "find" method, that returned "false".
void RBTree::insert(
    RBTreeNode* parentNode,
    RBTreeNodeValue* key
) {
    assert(parentNode != 0);
    int n = (-1);
    if (parentNode->value != 0)
        n = key->compareTo(*((const RBTreeNodeValue*) parentNode->value));
    assert(n != 0);
    RBTreeNode* x = new RBTreeNode();
    x->value = (void*) key;
    x->parent = parentNode;
    if (parentNode == &header)
        x->red = false;     // The root of tree is black
    else
        x->red = true;      // Paint the new node in red
    if (n < 0) {
        // Insert as a left son
        assert(parentNode->left == 0);
        parentNode->left = x;
    } else {
        // Insert as a right son
        assert(parentNode != &header && parentNode->right == 0);
        parentNode->right = x;
    }
    ++numNodes;

    if (x != root())
        rebalanceAfterInsert(x);
}

// Rotate a node x to the left    //
//        x                y      //
//       / \              / \     //
//      a   y    --->    x   c    //
//         / \          / \       //
//        b   c        a   b      //

void RBTree::rotateLeft(RBTreeNode* x) {
    RBTreeNode* y = x->right;
    assert(y != 0);
    RBTreeNode* p = x->parent;
    y->parent = p;
    if (x == p->left) {
        // x is the left son of its parent
        p->left = y;
    } else {
        // x is the right son of its parent
        p->right = y;
    }
    x->right = y->left;
    if (y->left != 0)
        y->left->parent = x;
    y->left = x;
    x->parent = y;
}

// Rotate a node x to the right   //
//        x                y      //
//       / \              / \     //
//      y   c    --->    a   x    //
//     / \                  / \   //
//    a   b                b   c  //

void RBTree::rotateRight(RBTreeNode* x) {
    RBTreeNode* y = x->left;
    assert(y != 0);
    RBTreeNode* p = x->parent;
    y->parent = p;
    if (x == p->left) {
        // x is the left son of its parent
        p->left = y;
    } else {
        // x is the right son of its parent
        p->right = y;
    }
    x->left = y->right;
    if (y->right != 0)
        y->right->parent = x;
    y->right = x;
    x->parent = y;
}

void RBTree::rebalanceAfterInsert(RBTreeNode* x) {
    assert(x->red);
    while(x != root() && x->parent->red) {
        if (x->parent == x->parent->parent->left) {
            // parent of x is a left son

            RBTreeNode* y = x->parent->parent->right; // y is the sibling of
                                             //            parent of x
            if (y != 0 && y->red) {          // if y is red
                x->parent->red = false;      //     color parent of x in black
                y->red = false;              //     color y in black
                x = x->parent->parent;       //     x = grandparent of x
                x->red = true;               //     color x in red
            } else {                         // else y is black
                if (x == x->parent->right) { //     if x is a right son
                    x = x->parent;           //         x = parent of x
                    rotateLeft(x);           //         left-rotate x
                }                            //     end if
                assert(x == x->parent->left);//     assert: x is a left son
                x->parent->red = false;      //     color parent of x in black
                x->parent->parent->red = true;  // color grandparent in red
                rotateRight(x->parent->parent); // right-rotate grandparent
            }                                // endif
        } else {
            // Mirror case: parent of x is a right son
            assert(x->parent == x->parent->parent->right);

            RBTreeNode* y = x->parent->parent->left; // y is the sibling of
                                             //           parent of x
            if (y != 0 && y->red) {          // if y is red
                x->parent->red = false;      //     color parent of x in black
                y->red = false;              //     color y in black
                x = x->parent->parent;       //     x = grandparent of x
                x->red = true;               //     color x in red
            } else {                         // else y is black
                if (x == x->parent->left) {  //     if x is a left son
                    x = x->parent;           //         x = parent of x
                    rotateRight(x);          //         right-rotate x
                }                            //     end if
                assert(x == x->parent->right); //   assert: x is a right son
                x->parent->red = false;      //     color parent of x in black
                x->parent->parent->red = true; //  color grandparent in red
                rotateLeft(x->parent->parent); //  left-rotate grandparent
            }                                // endif
        }
    } // end while

    // Always color the root in black
    if (x == root()) {
        x->red = false;
    }
}

void RBTree::eraseNode(RBTreeNode* node) {
    delete (RBTreeNodeValue*) node->value;
}

void RBTree::clear() {
    removeSubtree(root());
}

int RBTree::removeSubtree(RBTreeNode* subTreeRoot) {
    int numRemoved = 0;
    if (subTreeRoot == 0)
        return 0;
    if (subTreeRoot->left != 0)
        numRemoved += removeSubtree(subTreeRoot->left); // recursive call
    if (subTreeRoot->right != 0)
        numRemoved += removeSubtree(subTreeRoot->right); // recursive call

    if (subTreeRoot->parent->left == subTreeRoot)
        subTreeRoot->parent->left = 0;
    else
        subTreeRoot->parent->right = 0;

    eraseNode(subTreeRoot);
    delete subTreeRoot;
    ++numRemoved;
    --numNodes;

    assert(numNodes >= 0);

    return numRemoved;
}

const RBTreeNode* RBTree::minimalNode(
    const RBTreeNode* subTreeRoot /* = 0 */
) const {
    const RBTreeNode* x = subTreeRoot;
    if (x == 0)
        x = root();
    while (x != 0 && x->left != 0)
        x = x->left;
    return x;
}

const RBTreeNode* RBTree::maximalNode(
    const RBTreeNode* subTreeRoot /* = 0 */
) const {
    const RBTreeNode* x = subTreeRoot;
    if (x == 0)
        x = root();
    while (x != 0 && x->right != 0)
        x = x->right;
    return x;
}

const RBTreeNode* RBTree::nextNode(const RBTreeNode* node) const {
    assert(node != 0);
    if (node == &header)
        return minimalNode();

    if (node->right != 0) {
        return minimalNode(node->right);
    } else if (node == node->parent->left) { // node is a left son
        return node->parent;
    } else {                                 // node is a right son
        const RBTreeNode* x = node->parent;
        while (x == x->parent->right)        // while x is a right son
            x = x->parent;
        return x->parent;
    }
}

const RBTreeNode* RBTree::previousNode(const RBTreeNode* node) const {
    assert(node != 0);
    if (node == minimalNode())
        return &header;

    if (node->left != 0) {
        return maximalNode(node->left);
    } else if (node == node->parent->right) { // node is a right son
        return node->parent;
    } else {                                  // node is a left son
        const RBTreeNode* x = node->parent;
        while (x->parent != 0 && x == x->parent->left) // while x is a left son
            x = x->parent;
        if (x->parent != 0) {
            return x->parent;
        } else {
            assert(x == &header);
            return x;
        }
    }
}

//
// Set (Map) based on Red-Black Tree
//
#ifndef TREESET_H
#define TREESET_H

#include "RBTree.h"

// An ABSTRACT class representing a key in TreeSet
class TreeSetKey {
public:
    TreeSetKey() {}
    virtual ~TreeSetKey() {}    // virtual destructor

    // Abstract class has virtual methods unimplemented

    // The keys can be compared
    virtual int compareTo(const TreeSetKey&) const = 0;

    bool operator==(const TreeSetKey& k) const {
        return (compareTo(k) == 0);
    }
    bool operator!=(const TreeSetKey& k) const {
        return (compareTo(k) != 0);
    }
    bool operator<(const TreeSetKey& k) const {
        return (compareTo(k) < 0);
    }
    bool operator<=(const TreeSetKey& k) const {
        return (compareTo(k) <= 0);
    }
    bool operator>(const TreeSetKey& k) const {
        return (compareTo(k) > 0);
    }
    bool operator>=(const TreeSetKey& k) const {
        return (compareTo(k) >= 0);
    }

    // The virtual method "clone" must allocate a copy of the object
    // in the dynamic memory. In any derived class Foo it must
    // be always implemented in the following way:
    // virtual Foo* clone() const { return new Foo(*this); }
    //
    virtual TreeSetKey* clone() const = 0;
};

// An ABSTRACT class representing a value of a key in TreeSet
class TreeSetValue {
public:
    TreeSetValue() {}
    virtual ~TreeSetValue() {}
    virtual TreeSetValue* clone() const = 0;
};

//
// class TreeSet implements "Map" or "Dictionary".
// It stores the set of pairs: (key, value).
// All keys are unique (different pairs have different keys).
//
// The implementation is based on the Red-Black Tree
//
class TreeSet: protected RBTree {
public:
    class Pair: public RBTreeNodeValue {
    public:
        const TreeSetKey* key;
        TreeSetValue* value;

        Pair():
            RBTreeNodeValue(),
            key(0),
            value(0)
        {}
        Pair(const TreeSetKey* k, TreeSetValue* v):
            RBTreeNodeValue(),
            key(k),
            value(v)
        {}

        virtual int compareTo(const RBTreeNodeValue& p) const {
            return key->compareTo(
                *((const Pair&) p).key
            );
        }
    };

    // Add a pair (key, value) to the set
    void add(const TreeSetKey* k, const TreeSetValue* v = 0);

    void remove(const TreeSetKey* key);

    // Return a value of a key
    TreeSetValue* value(const TreeSetKey* k) const;

    TreeSetValue* operator[](const TreeSetKey* k) const {
        return value(k);
    }

    bool contains(const TreeSetKey* k) const;

    int size() { return RBTree::size(); }

    class const_iterator: public RBTree::const_iterator {
    public:
        const_iterator(const RBTree::const_iterator& i):
            RBTree::const_iterator(i)
        {}

        const Pair& operator*() const { // Dereference
            return *((const Pair*) node->value);
        }
        const Pair* operator->() const {
            return &(operator*());
        }
    };

    class iterator: public RBTree::iterator {
    public:
        iterator(const RBTree::iterator& i):
            RBTree::iterator(i)
        {}

        Pair& operator*() const { // Dereference
            return *((Pair*) node->value);
        }
        Pair* operator->() const {
            return &(operator*());
        }
    };

    const_iterator begin() const {
        return RBTree::begin();
    }
    const_iterator end() const {
        return RBTree::end();
    }

    iterator begin() {
        return RBTree::begin();
    }
    iterator end() {
        return RBTree::end();
    }
};

#endif

#include "TreeSet.h"

bool TreeSet::contains(const TreeSetKey* k) const {
    Pair p(k, 0);
    return find(&p);
}

void TreeSet::add(const TreeSetKey* k, const TreeSetValue* v /* = 0 */) {
    Pair key(k, 0);
    RBTreeNode* node;
    if (find(&key, root(), &node)) {
        // The key is already in the set
        Pair* p = (Pair*) node->value;
        delete p->value;            // Remove the old value
        if (v != 0) {
            p->value = v->clone();  // Assign the new value
        } else {
            p->value = 0;
        }
    } else {
        // Add the pair to the set
        Pair* p = new Pair(k->clone(), v->clone());
        insert(node, p);
    }
}

TreeSetValue* TreeSet::value(const TreeSetKey* k) const {
    Pair key(k, 0);
    RBTreeNode* node;
    if (find(&key, root(), &node)) {
        // The key is in the set
        Pair* p = (Pair*) node->value;
        return p->value;
    } else {
        // The key is not found
        return 0;
    }
}

//
// Test of class RBTree representing Red-Black Tree
// We work with the set of integer values
//
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <time.h>
#include "RBTree.h"

static bool writeIntegerTree(const RBTreeNode* root, FILE* f, int level = 0);
static bool readIntegerTree(RBTree& tree, FILE* f);
static void printHelp();

class Integer: public RBTreeNodeValue {
public:
    int number; // The value of a node is an integer number

    Integer(): 
        RBTreeNodeValue(), 
        number(0) 
    {}

    Integer(int n):
        RBTreeNodeValue(), 
        number(n)
    {}

    Integer(const Integer& n):
        RBTreeNodeValue(), 
        number(n.number)
    {}

    virtual int compareTo(const RBTreeNodeValue& v) const {
        return (number - ((const Integer&) v).number);
    }

    virtual ~Integer() {}
};

int main() {
    printf("Test of Red-Black Tree class\n");
    printHelp();
    FILE* f = NULL;

    // Define a random tree with 20 nodes
    time_t t = time(0); // Current time in seconds since 1 Jan 1970
    srand(t);

    RBTree tree;
    char line[256];
    while (true) {
        printf("Command>");
        if (fgets(line, 254, stdin) == NULL)
            break;

        // Parse a command
        line[254] = 0;
        int len = strlen(line);
        // Remove "\r\n" at the end of line
        if (len > 0 && line[len-1] == '\n') {
            line[len-1] = 0; --len;
        }
        if (len > 0 && line[len-1] == '\r') {
            line[len-1] = 0; --len;
        }

        int i = 0;
        // Skip a space in beginning of line
        while (i < len && isspace(line[i])) ++i;
        int commandBeg = i;
        while (i < len && isalpha(line[i])) ++i;
        int commandEnd = i;
        int commandLen = commandEnd - commandBeg;

        if (
            strncmp("generateTree", line+commandBeg, commandLen) == 0
        ) {
            while (i < len && isspace(line[i])) ++i; // Skip a space
            if (i >= len || !isdigit(line[i])) {
                printf("Incorrect command.\n");
                printHelp();
                continue;       // Incorrect command
            }
            int n = atoi(line + i);
            while (i < len && isdigit(line[i])) ++i; // Skip a number
            while (i < len && isspace(line[i])) ++i; // Skip a space
            if (i >= len) {
                printf("Incorrect command.\n");
                printHelp();
                continue;       // Incorrect command
            }

            if ((f = fopen(line+i, "w")) == NULL) {
                perror("Could not open a file for writing");
                continue;
            }

            // Generate a random tree
            tree.clear();
            int j = 0;
            if (n > 100)
                n = 100;        // Maximum 100 nodes
            while (j < n) {
                Integer num(rand() % 100 + 1);
                RBTreeNode* node;
                if (!tree.find(&num, tree.root(), &node)) {
                    Integer* v = new Integer(num);
                    tree.insert(node, v);
                    ++j;
                }
            }

            writeIntegerTree(tree.root(), f);
            fclose(f); f = NULL;

            // Print a tree to stdout
            writeIntegerTree(tree.root(), stdout);

        } else if (
            strncmp("readTree", line+commandBeg, commandLen) == 0
        ) {

            while (i < len && isspace(line[i])) ++i; // Skip a space

            if ((f = fopen(line+i, "r")) == NULL) {
                perror("Could not open a file for reading");
                continue;
            }

            // Read a tree from a file
            if (readIntegerTree(tree, f)) {
                writeIntegerTree(tree.root(), stdout);
            } else {
                printf("Incorrect format.\n");
            }
            fclose(f); f = NULL;

        } else if (
            strncmp("quit", line+commandBeg, commandLen) == 0
        ) {

            break;

        } // end if
    } // end while

    return 0;
}

// Recursive function
static bool writeIntegerTree(
    const RBTreeNode* root, FILE* f, int level /* = 0 */
) {
    int i;
    for (i = 0; i < level; ++i) {
        fprintf(f, "    ");
    }
    if (root == 0) {
        fprintf(f, "null\n");
    } else {
        Integer* v = (Integer*) root->value;
        fprintf(f, "%d ", v->number);
        if (root->red)
            fprintf(f, " red\n");
        else
            fprintf(f, " black\n");
        writeIntegerTree(root->left, f, level+1);
        writeIntegerTree(root->right, f, level+1);
    }

    return true;
}

// Read a tree from a stream
// Example:                               //
//                 10 black               //
//                /        \              //
//            5 red         15 black      //
//           /   \          /    \        //
//        null  7 black   null    null    //
//               /  \                     //
//            null  null                  //
// is represented as                      //
// 10 black
//     5 red
//         null
//         7 black
//             null
//             null
//     15 black
//         null
//         null
// This is a recursive function
//
static bool readIntegerTree(RBTree& tree, FILE* f) {
    tree.clear();

    char line[256];
    if (fgets(line, 254, f) == NULL)
        return false;
    line[254] = 0;
    int len = strlen(line);
    // Remove "\r\n" at the end of line
    if (len > 0 && line[len-1] == '\n')
        --len;
    if (len > 0 && line[len-1] == '\r')
        --len;
    int i = 0;
    while (i < len && isspace(line[i]))
        ++i;
    if (i < len && (isdigit(line[i]) || line[i] == '-')) {
        int n = atoi(line + i);
        bool red = false;
        // Skip a number
        while (i < len && (isdigit(line[i]) || line[i] == '-'))
            ++i;
        if (i >= len || !isspace(line[i]))
            return false;
        // Skip space
        while (i < len && isspace(line[i]))
            ++i;
        if (i < len && line[i] == 'r')
            red = true;

        RBTree leftSubtree;
        if (!readIntegerTree(leftSubtree, f))   // recursive call
            return false;

        RBTree rightSubtree;
        if (!readIntegerTree(rightSubtree, f))  // recursive call
            return false;

        RBTreeNode* rootNode = new RBTreeNode();
        rootNode->red = red;
        rootNode->value = new Integer(n);
        tree.header.left = rootNode;
        rootNode->parent = &(tree.header);
        tree.numNodes = 1;

        if (leftSubtree.size() > 0) {
            // Link root node with left subtree
            rootNode->left = leftSubtree.root();
            leftSubtree.root()->parent = rootNode;
            tree.numNodes += leftSubtree.size();

            // Erase left subtree
            leftSubtree.header.left = 0;
            leftSubtree.numNodes = 0;
        }

        if (rightSubtree.size() > 0) {
            // Link root node with right subtree
            rootNode->right = rightSubtree.root();
            rightSubtree.root()->parent = rootNode;
            tree.numNodes += rightSubtree.size();

            // Erase right subtree
            rightSubtree.header.left = 0;
            rightSubtree.numNodes = 0;
        }
    }

    return true;
}

static void printHelp() {
    printf(
        "Commands:\n"
        "  generateTree n fileName\t"
        "generate random tree with n nodes\n"
        "\t\t\t\tand write it to the file \"fileName\"\n\n"
        "  readTree fileName\t"
        "\tread a tree from the file \"fileName\"\n\n"
        "  quit\t\t\t\tquit\n"
    );
}

//
// Test of class TreeSet:
// Calculate the set of all words in a text,
// and for every word calculate a number of its inclusions
// in the text.
//
// The program reads a text either from a file
// or from standard input stream, depending on how it is called:
//
// 1) the command
//        wrdFreq input_file
//    reads the text from "input_file";
//
// 2) the command
//        wrdFreq
//    without arguments reads the text from standard input.
//    It allows, for instance, the following possibility (in Unix):
//    Define the set of words in all Unix standard include files.
//    It can be done by the command
//        cat /usr/include/*.h | ./wordfreq
//
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <ctype.h>
#include <assert.h>

#include "TreeSet.h"

static void printHelp();

//
// class Word represents a word in a text.
// This is the dynamic array of characters.
// The class Word is derived from class TreeSetKey,
// so the object of this class can be used
// as keys in TreeSet
//
// Interface of class Word:
//
class Word: public TreeSetKey {
    char* str;
    int   len;
    int capacity;
public:
    Word();
    Word(const char* s, int l = (-1));  // (-1) means "undefined"
    Word(const Word& w);
    virtual ~Word() { delete[] str; }
    virtual Word* clone() const { return new Word(*this); }

    int length() const { return len; }
    int size() const { return len; }
    void initialize();

    // Add a character to the end of word
    Word& operator+=(int c);

    // Convertor to C-string
    operator const char *() { return str; }
    const char* getString() const { return str; }

    virtual bool operator==(const TreeSetKey& s) const {
        return (strcmp(str, ((Word&) s).str) == 0);
    }

    virtual int compareTo(const TreeSetKey& key) const {
        return strcmp(str, ((const Word&) key).str);
    }
};

// Implementation of class Word
Word::Word():
    TreeSetKey(),
    str(0),
    len(0),
    capacity(0)
{}

Word::Word(const char* s, int l):
    TreeSetKey(),
    str(0),
    len(0),
    capacity(0)
{
    assert(str != 0);
    if (l < 0)
        l = strlen(str);
    len = l;
    capacity = len + 1;
    str = new char[capacity];
    memmove(str, s, len);
    str[len] = 0;
}

Word::Word(const Word& w):
    TreeSetKey(w),
    str(0),
    len(w.len),
    capacity(w.capacity)
{
    if (w.len > 0) {
        str = new char[w.capacity];
        memmove(str, w.str, capacity);
    }
}

void Word::initialize() {
    len = 0;
    if (capacity > 0)
        str[0] = 0;
}

// Add a character to the end of word
Word& Word::operator+=(int c) {
    if (capacity <= len+1) {
        // Allocate additional space
        int extent = capacity / 8;
        if (extent < 16)
            extent = 16;
        char* new_str = new char[capacity + extent];
        memmove(new_str, str, len);
        delete[] str;
        capacity += extent;
        str = new_str;
    }
    str[len] = (char) c;
    ++len;
    str[len] = 0;
    return *this;
}

// class Integer represents the number of inclusions of word in text
class Integer: public TreeSetValue {
public:
    int number;
public:
    Integer(): 
        TreeSetValue(),
        number(0) 
    {}

    Integer(const Integer& i): 
        TreeSetValue(),
        number(i.number) 
    {}

    Integer(int n): 
        TreeSetValue(),
        number(n) 
    {}

    virtual ~Integer() {}

    virtual Integer* clone() const {
        return new Integer(*this);
    }
};

int main(int argc, char *argv[]) {
    TreeSet set;
    FILE* input;
    if (argc > 1) {
        if (*argv[1] == '-') {
            printHelp();
            return 0;
        }
        input = fopen(argv[1], "r");
        if (input == 0) {
            perror("Cannot open an input file");
            return 1;
        }
    } else {
        input = stdin;  // Standard input
    }

    Word currentWord;
    bool wasLetter = false;
    bool endOfFileDetected = false;
    while (!endOfFileDetected) {
        int c = fgetc(input);
        if (isalpha(c)) {       // If this is an English letter
            currentWord += c;   // Add a letter to a word
            wasLetter = true;
        } else {
            if (wasLetter) {
                // The end of current word detected
                if (set.contains(¤tWord)) {
                    Integer* val = (Integer*) set.value(¤tWord);
                    ++(val->number);    // Increment a number of inclusions
                } else {
                    Integer unit(1);
                    set.add(¤tWord, &unit);
                }
                currentWord.initialize();
            }
            wasLetter = false;
            endOfFileDetected = (c == EOF);
        }
    }

    // Print the set of words in the text and
    // define the most frequent word
    printf("The text contains the following words:\n");

    TreeSet::const_iterator i = set.begin();
    TreeSet::const_iterator e = set.end();
    Word mostFrequentWord;
    int numberOfMostFrequentWords = 0;
    while (i != e) {
        const TreeSet::Pair& pair = *i;
        const Word* w = (const Word *) pair.key;
        int number = ((const Integer *) pair.value)->number;
        printf(
            "%d\t%s\n", number, w->getString()
        );

        // Define the most frequent word
        if (number > numberOfMostFrequentWords) {
            numberOfMostFrequentWords = number;
            mostFrequentWord = *w;
        }
        ++i;
    }

    printf(
        "----\n"
        "Number of different words in the text = %d\n",
        set.size()
    );

    printf(
        "The most frequent word is \"%s\", included %d times.\n",
        mostFrequentWord.getString(), numberOfMostFrequentWords
    );

    return 0;
}

static void printHelp() {
    printf(
        "Calculate the set of all words in a text,\n"
        "and for every word calculate a number of its inclusions\n"
        "in the text.\n"
        "Usage:\n"
        "    wordfreq [input_file]\n"
        "The program reads a text either from a file\n"
        "or from standard input stream, depending on how it is called.\n\n"
        "EXAMPLES:\n"
        "1) the command\n"
        "       ./wordfreq input_file\n"
        "   reads the text from \"input_file\";\n"
        "\n"
        "2) the Unix-command\n"
        "       cat /usr/include/*.h | ./wordfreq\n"
        "   defines the set of words in all standard include files.\n"
    );
}

uses crt;

type
  string_20 = string[20]; { Это - тип данных, хранимых в КЧД }

  { Узел дерева }
  PTreeNode = ^TreeNode;
  TreeNode = object
  public
    constructor init(_s: string; _parent: PTreeNode);

  public
    Color: integer; { Цвет листа: 0 = Черный, 1 = Красный }
    left, right, parent,
    duplicates: PTreeNode; { Список дубликатов (не стал удалять, иногда это нужно) }

    data_string: string_20;
    deleted: boolean;      { Флаг для "ленивого" удаления }
  end;

  RBTree = object
  public
    constructor init;
    destructor done;

    function Search(s: string): integer;
    { Возвращает +1, если строка присутствует в дереве, и -1 если ее там нет }
    
    function SearchFirstMatch(s: string): PTreeNode;
    { Работает точно так же, как и Search, но возвращает указатель
       типа PTreeNode на первый подходящий элемент }
       
    procedure Insert(s: string);
    { Добавляет новый элемент в дерево }
    function InsertItem(s: string; node: PTreeNode): PTreeNode;

    function Remove(s: string): boolean;
    { Удаляет заданную строку }

    procedure LoadFromFile(fn: string); { Загружает дерево из текстового файла 
                                             (не реализовано) }
    procedure SaveToFile(var f: text);  { Сохраняет дерево в текстовый файл }

    function LeftDepth: integer;  { Находит глубину левого поддерева }
    function RightDepth: integer; { Находит глубину правого поддерева }
    function NumberOfLeaves(p: PTreeNode): integer;
    { Находит число листьев в дереве }

  public
    root: PTreeNode;

  private
    procedure LeftRotation(node: PTreeNode);
    procedure RightRotation(node: PTreeNode);
    { "Ротации" (повороты), используемые при вставке для балансировки дерева }
    function TreeDepth(p: PTreeNode): integer;
    { Рекурсивная функция для нахождения глубины дерева, с корнем в P }

    procedure DeleteTree(p: PTreeNode);
    { Деструктор вызывает эту процедуру для удаления дерева }
    procedure SaveNode(level: integer;
              const node: PTreeNode; var f: text);
    { Рекурсивная процедура сохранения узла в тестовый файл F }
  end;


Const
  { Цвета для узлов }
  node_Red   = 1;
  node_Black = 0;

constructor TreeNode.init(_s: string; _parent: PTreeNode);
begin
  data_string := _s;
  left := nil; right := nil;
  parent := _parent; { Указатель на предка }

  Duplicates := nil; { Изначально у узла нет дубликатов }
  Color := node_Red; { Новые узлы становятся "Красными" }
  Deleted := False;  { Этот узел не удален }
end;

{
  Функция сравнения строк (ведомые пробелы не принимаются во внимание)
}
function compare(s1, s2: string): integer;

  procedure trim(var s: string);
  begin
    while s[length(s)] = ' ' do delete(s, length(s), 1);
  end;

begin

  trim(s1); trim(s2);
  if s1 < s2 then compare := -1
  else if s1 > s2 then compare := +1
       else compare := 0;

end;

constructor RBTree.init;
begin
  root := nil;
end;
destructor RBTree.done;
begin
  DeleteTree(Root);
  { DeleteTree освобождает динамическую память }
end;


procedure RBTree.DeleteTree(p: PTreeNode);
begin
  if p <> nil then begin
    DeleteTree(p^.Left);  { Удалить левое поддерево }
    DeleteTree(p^.Right); { Удалить правое поддерево }
    DeleteTree(p^.Duplicates);
    dispose(p);
  end;
  p := nil; { Узел более не используется }
end;

{
  При вставке элемента могут произойти 3 случая:
  1. Новый узел и его "дядя" - "Красные"
  2. Новый узел красный, "дядя" - "Черный", и узел - левый потомок
  3. Новый узел красный, "дядя" - "Черный", и узел - правый потомок
}
procedure RBTree.Insert(s: string);
{ Создает новый узел для хранения строки }
var
  node, node_2: PTreeNode;
begin
  node := InsertItem(s, root); { Вставить строку в дерево }

  if node = nil then exit; { Изменять дерево не нужно }

  while(node <> root) and (node^.parent^.color = node_Red) do begin

    { Проверяем, находится ли узел в левом поддереве }
    if node^.parent = node^.parent^.parent^.left then begin

      node_2 := node^.parent^.parent^.right; { Делаем node2 "дядей" нашего узла }
      
      { Если "дядя" красного цвета - это случай 1 }
      if (node_2 <> nil) and (node_2^.Color = node_Red) then begin
        node^.parent^.Color := node_Black; { Изменяем цвет "родителя" на черный }
        node_2^.Color := node_Black;       { Изменяем цвет "дяди" на черный }
        node^.Parent^.Parent^.Color := node_Red; { Делаем "дедушку" красным }
        node := node^.Parent^.Parent;            { Двигаемся к вершине дерева для 
                                                   проведения дополнительных 
                                                   исправлений }
      end
      else begin { "дядя" - черного цвета, случай 2 или 3 }
      
        if Node = Node^.Parent^.Right then begin { Проверяем на случай №3 }
          Node := Node^.Parent; { Узел - правый потомок, это как раз случай №3... }
          LeftRotation(Node);   { ... который требует левую "ротацию" }
        end;
        Node^.Parent^.Color := node_Black;       { Установка для случаев №2... }
        Node^.Parent^.Parent^.Color := node_Red; { ... и №3 }
        RightRotation(Node^.Parent^.Parent);
        
      end;

    end
    else begin { узел - в правом поддереве }

      node_2 := Node^.Parent^.Parent^.Left; { Делаем node2 "дядей" нашего узла }
      { Если "дядя" красного цвета - это случай 1 }
      if(node_2 <> nil) and (node_2^.Color = node_Red) then begin
        Node^.Parent^.Color := node_Black;       { Изменяем цвет "родителя" на черный }
        Node_2^.Color := node_Black;             { Изменяем цвет "дяди" на черный }
        Node^.Parent^.Parent^.Color := node_Red; { Делаем "дедушку" красным }
        Node := Node^.Parent^.Parent;            { Двигаемся к вершине дерева... }
      end
      else begin { "дядя" - "черный", случай №2 или №3 }

        { Проверяем на случай №3 ("лево" и "право" обменяны местами) }
        if Node = Node^.Parent^.Left then begin
          Node := Node^.Parent; { Узел - левый потомок, это как раз случай №3... }
          RightRotation(Node);  { ... который требует правую "ротацию" }
        end;
        Node^.Parent^.Color := node_Black;       { Установка для случаев №2... }
        Node^.Parent^.Parent^.Color := node_Red; { ... и №3 }
        LeftRotation(Node^.Parent^.Parent);
      end;

    end

  end;

  { По правилу КЧД корень дерева должен быть черным }
  Root^.Color := node_Black;
end;


function RBTree.InsertItem(s: string; node: PTreeNode): PTreeNode;
var
  comparison: integer;
  GreaterThanLeft, LessThanRight: boolean;
  T: PTreeNode;
begin

  if root = nil then begin

    root := new(PTreeNode, init(s, nil)); { устанавливаем корень }
    
    { По правилу КЧД корень дерева должен быть черным }
    root^.Color := node_Black;
    InsertItem := root; exit

  end;

  while True do begin

    comparison := compare(s, node^.data_string);

    if node^.Deleted then begin
      { Для начала проверим, является ли узел "удаленным".
        Если это так, то существует возможность использовать "удаленный" узел
        для хранения новой записи, если она должна будет находиться
        между двумя "потомками" }
      if node^.Left = nil then GreaterThanLeft := true
      else
        { (В случае, если compare() < 0): строка не больше чем левый потомок,
          поэтому "удаленный" узел не может использоваться для хранения новой записи }
        GreaterThanLeft := (compare(s, node^.left^.data_string) > 0);

      if node^.Right = nil then LessThanRight := true
      else
        { (В случае, если compare() < 0): строка не больше чем правый потомок,
          поэтому "удаленный" узел не может использоваться для хранения новой записи }
        LessThanRight := (compare(s, node^.right^.data_string) > 0);

      if GreaterThanLeft and LessThanRight then begin
        { "удаленный" узел может использоваться для хранения новой записи }
        node^.data_string := s;
        node^.Deleted := false; { удел больше "удаленным" не считать }
        InsertItem := nil; exit

        {
        возвращаем NIL, чтобы избежать "ротаций" дерева, т.к. элемент, значение
        которого было изменено, находится на своем месте
        }
      end;

    end;

    if comparison < 0 then begin
      { Если Left пусто, помещаем новый узел туда }
      if Node^.Left = nil then begin
        Node^.Left := new(PTreeNode, init(s, Node)); { Добавляем новый узел ... }
        InsertItem := Node^.Left;                    { ... как левого потомка }
        exit
      end
      else Node := Node^.Left; { Проверить левое поддерево }
    end

    else
      if comparison > 0 then begin
        { Если Right пусто, помещаем новый узел туда }
        if Node^.Right = nil then begin
          Node^.Right := new(PTreeNode, init(s, Node)); { Добавляем новый узел ... }
          InsertItem := Node^.Right;                    { ... как правого потомка }
          exit
        end
        else Node := Node^.Right; { Проверить правое поддерево }
      end
      else begin { узел - дубликат }

        T := node;
        { находим конец списка дубликатов }
        while(T^.Duplicates <> nil) do T := T^.Duplicates;
        T^.Duplicates := new(PTreeNode, init(s, T));
        InsertItem := nil; exit
        {
        возвращаем NIL, чтобы избежать "ротаций" дерева, т.к.
        мы просто изменили список дубликатов
        }
      end;
  end;

end;


function RBTree.Remove(s: string): boolean;
var T, prev_node, node: PTreeNode;
begin
  Remove := False;
  Node := SearchFirstMatch(s); { Найдем подходящий узел в дереве }
  if node = nil then exit;     { Строка не была найдена - выход }

  if node^.Duplicates <> nil then begin

    {
      если есть дубликаты - то один из дубликатов может занять
      место удаляемой записи
    }
    T := node;
    while T^.Duplicates <> nil do begin
      prev_node := T;
      T := T^.Duplicates;
    end;
    node^.data_string := T^.data_string;
    { Копируем содержимое последнего дубликата в ту запись, которую будем удалять }
    dispose(T);
    prev_node^.Duplicates := nil;
    { "отсекаем" последний элемент списка дубликатов }

    Remove := true; { удаление было успешным }
  end
  else
    Node^.Deleted := true; { Помечаем узел как "удаленный" для "ленивого" удаления }
  Remove := True
end;


function RBTree.Search(s: string): integer;
var
  node: PTreeNode;
  comparison: integer;
begin
  Search := -1;

  node := root;
  while Node <> nil do begin

    comparison := compare(s, node^.data_string);
    if comparison < 0 then Node := Node^.Left { просматриваем левое поддерево }
    else
      if comparison < 0 then Node := Node^.Right { просматриваем правое поддерево }
      else
        if Node^.Deleted then exit
          { если узел помечен на удаление - то не принимать его во внимание, выход }
        else begin
          { Строка найдена }
          search := 1; exit
        end;
  end;
  { Запись не найдена }
end;

function RBTree.SearchFirstMatch(s: string): PTreeNode;
{ Возвращает указатель на первый узел, хранящий заданную строку }
var
  node: PTreeNode;
  comparison: integer;
begin
  SearchFirstMatch := nil;

  node := root;
  while Node <> nil do begin

    comparison := compare(s, node^.data_string);
    if comparison < 0 then Node := Node^.Left { просматриваем левое поддерево }
    else
      if comparison > 0 then Node := Node^.Right { просматриваем правое поддерево }
      else
        if Node^.Deleted then exit
          { если узел помечен на удаление - то не принимать его во внимание, выход }
        else begin
          { Строка найдена }
          SearchFirstMatch := node; exit
        end;

  end;
end;

procedure RBTree.SaveToFile(var f: text);
{
  сохраняет узлы в Прямом (нисходящем) порядке
}
begin
  { Вызываем рекурсию }
  SaveNode(0, root, f);
end;

procedure RBTree.SaveNode(level: integer;
          const node: PTreeNode; var f: text);
const
  _color: array[0 .. 1] of char = ('B', 'R');
begin
  if node <> nil then begin

    if not node^.Deleted then begin
      writeln(f, '':3*level,
              node^.data_string + ' ('+_color[node^.Color]+')');
    end;
    SaveNode(level + 1, node^.Left, f);  { Сохраним узлы левого поддерева }
    SaveNode(level + 1, node^.Right, f); { Сохраним узлы правого поддерева }

  end;
end;

procedure RBTree.LoadFromFile(fn: string);
begin
  (* Не реализовано *)
end;

function RBTree.LeftDepth: integer;
begin
  LeftDepth := TreeDepth(Root^.Left); { Измеряем левое поддерево }
end;

function RBTree.RightDepth: integer;
begin
  RightDepth := TreeDepth(Root^.Right); { Измеряем правое поддерево  }
end;

function RBTree.TreeDepth(p: PTreeNode): integer;
var _left, _right: integer;
begin
  _left := 0; _right := 0;
  if p^.Left <> nil then
    _left := TreeDepth(p^.Left);   { Взять глубину левого поддерева }
  if p^.Right <> nil then
    _right := TreeDepth(p^.Right); { Взять глубину правого поддерева }

  if _left > _right then     { проверяем, какое поддерево "глубже" }
    TreeDepth := _left + 1   { вернем глубину левого поддерева + 1 }
  else
    TreeDepth := _right + 1; { вернем глубину правого поддерева + 1 }
end;


function RBTree.NumberOfLeaves(p: PTreeNode): integer;
var total: integer;
begin
  NumberOfLeaves := 1;

  total := 0;
  if (p^.Left = nil) and (p^.Right = nil) then exit; { узел является "листом" }

  { считаем число листьев в левом поддереве }
  if p^.Left <> nil then inc(total, NumberOfLeaves(p^.Left));
  { считаем число листьев в правом поддереве }
  if p^.Right <> nil then inc(total, NumberOfLeaves(p^.Right));

  NumberOfLeaves := total;
  { и возвращаем общее количество листьев в дереве }
end;

procedure RBTree.LeftRotation(node: PTreeNode);
var Right: PTreeNode;
begin
  Right := node^.Right; { hold node's right child }

  { make the node's right child its right child's left child }
  node^.Right := Right^.Left;
  if Right^.Left <> nil then
    Right^.Left^.Parent := Node; { point the child to its new parent }

  if Right <> nil then
    Right^.Parent := Node^.Parent; { point the child to its new parent }

  if Node^.Parent <> nil then begin { if node is not the root }
    if Node = Node^.Parent^.Left then { if node is a left child }
      Node^.Parent^.Left := Right   { make node's right child its 
                                                         parent's left child }
    else
      Node^.Parent^.Right := Right; { make node's right child its 
                                                         parent's right child }
  end
  else
    Root := Right; { node's right child is now the root }

  Right^.Left := Node; { node becomes its right child's left child }
  if Node <> nil then
    Node^.Parent := Right; { point node to its new parent }
end;

procedure RBTree.RightRotation(node: PTreeNode);
var Left: PTreeNode;
begin
  Left := node^.Left; { hold node's left child }

  { make the node's left child its left child's right child }
  Node^.Left := Left^.Right;
  if Left^.Right <> nil then
    Left^.Right^.Parent := Node; { point the child to its new parent }

  if Left <> nil then
    Left^.Parent := Node^.Parent; { point the child to its new parent }

  if Node^.Parent <> nil then begin { if node is not the root }
    if Node = Node^.Parent^.Right then { if node is a right child }
      Node^.Parent^.Right := Left { make node's left child its parent's right child }
    else
      Node^.Parent^.Left := Left; { make node's left child its parent's left child }
  end
  else
    Root := Left; { node's left child is now the root }

  Left^.Right := Node; { node becomes its left child's right child }
  if Node <> nil then
    Node^.Parent := Left; { point node to its new parent }
end;


{ Собственно, программа, иллюстрирующая использование КЧД }
var
  console: text;
  s: string_20;
  tree: RBTree;

begin
  assigncrt(console);
  rewrite(console);
  tree.init;

  {
    Вводим следующую последовательность строк:
    one
    two
    three
    four
    five
  }
  repeat
    write('enter new string (20 chars max): '); readln(s);
    if s <> '' then begin
      tree.insert(s);
      Writeln('**');
        tree.SaveToFile(console); { Выводим дерево на консоль }
      writeln('**');
    end;
  until s = '';
  tree.SaveToFile(console);

  { Проверяем работу Search }
  if tree.search('four') = 1 then writeln('found')
  else writeln('not found');

  { Проверяем работу Remove }
  tree.Remove('four');
  tree.SaveToFile(console);

  tree.done;
  close(console);
end.