Binary tree iterator

bintree_iterator.cpp

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

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

class InorderIterator {
private:
  TreeNode *root;
  TreeNode *curent;
  stack<TreeNode *> stk;

public:
  InorderIterator(TreeNode *_root) {
    root = _root;
    curent = _root;
  }

  bool hasNext() {
    return (curent != NULL || !stk.empty());
  }

  TreeNode *next() {
    while (curent != NULL || !stk.empty()) {
      if (curent != NULL) {
        stk.push(curent);
        curent = curent->left;
      } else { // !stk.empty()
        TreeNode *ret = stk.top();
        stk.pop();
        curent = ret->right;
        return ret;
      }
    }
  }
};

/*
void inorderRe(TreeNode *node) {
  if (node == NULL) return;
  inorderRe(node->left);
  cout << node->val << ", ";
  inorderRe(node->right);
}

void printTreeInorder(TreeNode* root) {
  inorderRe(root);
  cout << endl;
}

void preorderRe(TreeNode *node) {
  if (node == NULL) return;
  cout << node->val << ", ";
  preorderRe(node->left);
  preorderRe(node->right);
}

void printTreePreorder(TreeNode* root) {
  preorderRe(root);
  cout << endl;
}
*/

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

*/

int main() {
  // build tree
  TreeNode *root = new TreeNode(1);
  root->left = new TreeNode(2);
  root->right = new TreeNode(3);
  TreeNode *node4 = new TreeNode(4);
  root->left->left = node4;
  node4->right = new TreeNode(7);
  root->right->left = new TreeNode(5);
  TreeNode *node6 = new TreeNode(6);
  root->right->right = node6;
  node6->right =  new TreeNode(8);

  // iterator tree
  InorderIterator iter(root);
  while (iter.hasNext()) {
    TreeNode *node = iter.next();
    cout << node->val << ", ";
  }
  cout << endl;
  return 0;
}