find the node just larger than the specified node in BST

NextInBST.cpp

#include <iostream>
#include <vector>
#include <stack>
using namespace std;

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

// left: <=; right: >
TreeNode* next(TreeNode* root, TreeNode* node) {
  TreeNode* t = root;
  stack<TreeNode*> stk;
  while(t->val != node->val) {
    stk.push(t);
    if (t->val > node->val) {
      t = t->left;
    } else {
      t = t->right;
    }
  }

  // it's either in the right sub-tree
  if (t->right != NULL) {
    t = t->right;
    while (t->left != NULL)
      t = t->left;
    return t;
  }

  // or it's in the parent path as itself is in left sub-tree
  TreeNode *prev = NULL;
  while (!stk.empty()) {
    prev = t;
    t = stk.top();
    stk.pop();
    if (t->left == prev) return t;
  }
  // it's the last node already
  return NULL;
}

void test(TreeNode *root, TreeNode *node, TreeNode *target) {
  TreeNode * n = next(root, node);
  if (n != target) {
    cout << "ERROR: " << node->val << " vs. " << n ->val << endl;
  }
}

/*
         5
      /    \
     2       9
    / \     /
  1   4    7
     /    / \
    3    6   8
*/

int main() {
  // build tree
  vector<TreeNode *> nodes(10);
  for (int i = 0; i < 10; i++) {
    nodes[i] = new TreeNode(i);
  }

  TreeNode* root = nodes[5];
  nodes[5]->left = nodes[2];
  nodes[2]->left = nodes[1];
  nodes[2]->right = nodes[4];
  nodes[4]->left = nodes[3];
  nodes[5]->right = nodes[9];
  nodes[9]->left = nodes[7];
  nodes[7]->left = nodes[6];
  nodes[7]->right = nodes[8];

  for (int i = 1; i <= 8; i++) {
    test(root, nodes[i], nodes[i+1]);
  }

  test(root, nodes[9], NULL);

  return 0;
}