BSTree изменить код

#ifndef BSTSET_H_
#define BSTSET_H_
 
#include <string>
#include <iostream>
 
using namespace std;
 
enum Types{
    INNER,
    LEAF
};
 
class BSTreeException{
    string text;
public:
    BSTreeException(string _text){
        text = _text;
    }
    BSTreeException(const char* _text){
        text = string(_text);
    }
    string &getError(){
        return text;
    }
};
 
template <class Type> class Node{
public:
    virtual int getType() = 0;
    virtual Node *add(Type value) = 0;
    virtual Node *remove(Type value) = 0;
    virtual Type &leftMax() = 0;
    virtual Type &rightMax() = 0;
    virtual int count(Type value) = 0;
    virtual int depth(Type value) = 0;
    virtual ~Node(){};
};
 
template <class Type> class InnerNode;
 
template <class Type> class LeafNode: public Node<Type>{
    int _count;
public:
    Type element;
    LeafNode(Type elem){
        element = elem;
        _count = 1;
    }
    virtual Node<Type> *add(Type value){
        if(value == element){
            _count ++;
        } else {
            if(value < element){
                return new InnerNode<Type>(new LeafNode<Type>(value), this);
            } else {
                return new InnerNode<Type>(this, new LeafNode<Type>(value));
            }
        }
        return this;
    }
 
 
// метод который нужно изменить
    virtual Node<Type> *remove(Type value){
        if(value == element){
            if (_count > 1) {
                _count --;
                return this;
            } else {
                delete this;
                return NULL;
            }
        }
        return this;
    }
    virtual int getType(){
        return LEAF;
    }
    virtual Type &leftMax(){
        return element;
    }
    virtual Type &rightMax(){
        return element;
    }
    virtual int count(Type value){
        if(value != element){
            throw new BSTreeException("Element not found.");
        } else {
            return _count;
        }
    }
    virtual int depth(Type value){
        if(value != element){
            throw new BSTreeException("Element not found.");
        } else {
            return 1;
        }
    }
};
 
template <class Type> class InnerNode: public Node<Type> {
    Type _rightMax;
    Type _leftMax;
public:
    Node<Type> *right;
    Node<Type> *left;
    InnerNode(LeafNode<Type> *_left, LeafNode<Type> *_right){
        right = (Node<Type> *) _right;
        left = (Node<Type> *) _left;
        _rightMax = _right->rightMax();
        _leftMax = _left->leftMax();
    }
 
         // метод который нужно изменить
    virtual Node<Type> *remove(Type value){
        if(value < right->leftMax()){
            left = left->remove(value);
            if (left == NULL) {
                Node<Type> *tempNode = right;
                right = NULL;
                delete this;
                return tempNode;
            }
            _leftMax = left->leftMax();
        } else {
            right = right->remove(value);
            if (right == NULL) {
                Node<Type> *tempNode = left;
                left = NULL;
                delete this;
                return tempNode;
            }
            _rightMax = right->rightMax();
        }
        return this;
    }
    virtual Node<Type> *add(Type value){
        if(value < left->leftMax()){
            _leftMax = value;
        } else if(value > right->rightMax()){
            _rightMax = value;
        }
        if(value < right->leftMax()){
            left = left->add(value);
        } else {
            right = right->add(value);
        }
        return this;
    }
    virtual int getType(){
        return INNER;
    }
    virtual Type &leftMax(){
        return _leftMax;
    }
    virtual Type &rightMax(){
        return _rightMax;
    }
    virtual int count(Type value){
        if(value < right->leftMax()){
            return left->count(value);
        } else {
            return right->count(value);
        }
    }
    virtual int depth(Type value){
        if(value < right->leftMax()){
            return left->depth(value) + 1;
        } else {
            return right->depth(value) + 1;
        }
    }
    ~InnerNode(){
        delete right;
        delete left;
    }
};
 
template <class Type> class BSTSET {
    Node<Type> *root;
public:
    BSTSET(){
        root = NULL;
    }
    void operator>>(Type element){
        if(root != NULL){
            root = root->remove(element);
        }
    }
    void operator<<(Type element){
        if(root != NULL){
            root = root->add(element);
        } else {
            root = new LeafNode<Type>(element);
        }
 
    }
    int operator()(Type element){
        int count;
        try{
            if (root != NULL) {
                count = root->count(element);
            } else {
                count = 0;
            }
        } catch (BSTreeException *e){
            count = 0;
            delete e;
        }
        return count;
 
    }
    int operator[](Type element){
        int depth;
        try{
            if (root != NULL) {
                depth = root->depth(element);
            } else {
                depth = 0;
            }
        } catch (BSTreeException *e){
            depth = -1;
            delete e;
        }
        return depth;
    }
    virtual ~BSTSET(){
        if(root != NULL){
            delete root;
        }
    };
    template<class OtherType> friend ostream &operator<<(ostream &out, BSTSET<OtherType> &tree);
};
 
template<class Type> void printTree(ostream &out, Node<Type> *node){
    if(node->getType() == LEAF){
        out << ((LeafNode<Type> *) node)->element << endl;
    } else {
        printTree<Type>(out, ((InnerNode<Type> *) node)->left);
        printTree<Type>(out, ((InnerNode<Type> *) node)->right);
    }
}
 
template<class Type> ostream &operator<<(ostream &out, BSTSET<Type> &tree){
    printTree<Type>(out, tree.root);
    return out;
}
 
#endif /* BSTSET_H_ */