find cloest k node in bst

bst_iterator.cpp

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <stack>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <climits>

using namespace std;

struct TreeNode {
    int val;
    TreeNode  *left, *right;
    TreeNode(int v): val(v), left(NULL), right(NULL) {}
};

struct Iterator {
    virtual bool hasNext() = 0;
    virtual TreeNode *next() = 0;
};

struct NextIterator : Iterator{
    stack<TreeNode*> s;
    NextIterator(TreeNode *root, int val) {
        while (root) {
            s.push(root);
            if (root->val < val) {
                root = root->right;
            } else if (root->val > val) {
                root = root->left;
            } else {
                break;
            }
        }
    }

    bool hasNext() {
        return !s.empty(); 
    }

    TreeNode *next() {
        if (!hasNext()) {
            return NULL;
        }
        TreeNode *node = s.top();
        TreeNode *ret = node;
        s.pop();
        if (node->right) {
            node = node->right;
            while (node) {
                s.push(node);
                node = node->left;
            }
        } else {
            TreeNode *parent;
            while (!s.empty()) {
                parent = s.top();
                if (parent->left == node) {
                    break;
                }
                node = parent;
                s.pop();
            }
        }
        return ret;
    }
};

struct PrevIterator : Iterator {
    stack<TreeNode*> s;
    PrevIterator(TreeNode *root, int val) {
        while (root) {
            s.push(root);
            if (root->val < val) {
                root = root->right;
            } else if (root->val > val) {
                root = root->left;
            } else {
                break;
            }
        }
    }

    bool hasNext() {
        return !s.empty(); 
    }

    TreeNode *next() {
        if (!hasNext()) {
            return NULL;
        }
        TreeNode *node = s.top();
        TreeNode *ret = node;
        s.pop();
        if (node->left) {
            node = node->left;
            while (node) {
                s.push(node);
                node = node->right;
            }
        } else {
            TreeNode *parent;
            while (!s.empty()) {
                parent = s.top();
                if (parent->right == node) {
                    break;
                }
                node = parent;
                s.pop();
            }
        }
        return ret;
    }
};

void cloest2(Iterator *it, int &m1, int &m2, int val) {
    for (int i = 0; i < 2 && it->hasNext(); ++i) {
        int prev = it->next()->val;
        if (abs(m1-val) > abs(prev-val)) {
            m2 = m1;
            m1 = prev;
        } else if (abs(m2-val) > abs(prev-val)) {
            m2 = prev;
        }
    }
}

void cloest2Val(TreeNode *root, int val) {
    Iterator *nextIt = new NextIterator(root, val);
    Iterator *prevIt = new PrevIterator(root, val);
    int m1 = INT_MAX, m2 = INT_MAX;
    m1 = nextIt->next()->val;
    prevIt->next();
    cloest2(prevIt, m1, m2, val);
    cloest2(nextIt, m1, m2, val);
    cout<<m1<<":"<<m2<<endl;
}

int main() {
    TreeNode *root = new TreeNode(8);
    root->left        = new TreeNode(4);
    root->right       = new TreeNode(13);
    root->left->left  = new TreeNode(2);
    root->left->right = new TreeNode(6); 
    cloest2Val(root, 9);

}