Sunday, July 19, 2015

LeetCode OJ - Binary Search Tree Iterator

Problem:

Please find the problem here.

Solution:

Using a stack to keep track of unvisited parents, we can do the traversal.

First of all, the first node that should be visited is the left most node, so we go down the left path and push all the node to the stack.

Next, when we visit a node, we pop them out from the unvisited parents. To find the next node, there are two cases, it is possible that the current visiting node has a right sub-tree, in that case, it would be the left-most node of the right sub-tree, so we will just do what we did. Otherwise, we find the last unvisited parent and try its right sub-tree again.

Code:

#include "stdafx.h"

// https://leetcode.com/problems/binary-search-tree-iterator/

#include "LEET_BINARY_SEARCH_TREE_ITERATOR.h"
#include <map>
#include <iostream>
#include <sstream>
#include <vector>
#include <string>
#include <stack>

using namespace std;

namespace _LEET_BINARY_SEARCH_TREE_ITERATOR
{
    struct TreeNode
    {
        int val;
        TreeNode *left;
        TreeNode *right;
        TreeNode(int x) : val(x), left(NULL), right(NULL) {}
    };

    class BSTIterator
    {
    private:
        stack<TreeNode *> unvisited_parents;
        TreeNode* will_visit;
    public:
        BSTIterator(TreeNode *root)
        {
            if (root != NULL)
            {
                while (root != NULL)
                {
                    this->unvisited_parents.push(root);
                    root = root->left;
                }
                this->will_visit = this->unvisited_parents.top();
            }
            else
            {
                this->will_visit = NULL;
            }
        }

        /** @return whether we have a next smallest number */
        bool hasNext()
        {
            return this->will_visit != NULL;
        }

        /** @return the next smallest number */
        int next()
        {
            TreeNode* should_return = this->will_visit;
            TreeNode* current = this->will_visit;
            this->unvisited_parents.pop();
            if (current->right != NULL)
            {
                current = current->right;
                while (current != NULL)
                {
                    this->unvisited_parents.push(current);
                    current = current->left;
                }
            }
            this->will_visit = this->unvisited_parents.size() == 0 ? NULL : this->unvisited_parents.top();

            return should_return->val;
        }
    };
};

using namespace _LEET_BINARY_SEARCH_TREE_ITERATOR;

int LEET_BINARY_SEARCH_TREE_ITERATOR()
{
    TreeNode a(3);
    TreeNode b(5);
    TreeNode c(1);
    TreeNode d(6);
    TreeNode e(2);
    TreeNode f(0);
    TreeNode g(8);
    TreeNode h(7);
    TreeNode i(4);
    a.left = &b;
    a.right = &c;
    b.left = &d;
    b.right = &e;
    c.left = &f;
    c.right = &g;
    e.left = &h;
    e.right = &i;
    BSTIterator iter(&a);
    while (iter.hasNext())
    {
        cout << iter.next() << ",";
    }
    cout << endl;
    return 0;
}

No comments :

Post a Comment