[leetcode]Binary Search Tree Iterator

使用栈来记录可能的路径,栈顶一直是下一个元素。

class BSTIterator {
public:
    stack<TreeNode *> path;
    
    BSTIterator(TreeNode *root) {
        path = stack<TreeNode *>();
        TreeNode *node = root;
        while (node != NULL) {
            path.push(node);
            node = node->left;
        }
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return !path.empty();
    }

    /** @return the next smallest number */
    int next() {
        TreeNode *node = path.top();
        path.pop();
        int ret = node->val;
        if (node->right != NULL) {
            node = node->right;
            path.push(node);
            while (node->left != NULL) {
                node = node->left;
                path.push(node);
            }
        }
        return ret;
    }
};

Python3

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class BSTIterator:

    def __init__(self, root: TreeNode):
        self.stack = [] # append, pop
        self.node = root
        while self.node is not None:
            self.stack.append(self.node)
            self.node = self.node.left
            

    def next(self) -> int:
        """
        @return the next smallest number
        """
        if len(self.stack) > 0:
            self.node = self.stack.pop()
            result = self.node.val
            self.node = self.node.right
            while self.node:
                self.stack.append(self.node)
                self.node = self.node.left
            return result

    def hasNext(self) -> bool:
        """
        @return whether we have a next smallest number
        """
        if len(self.stack):
            return True
        return False
        


# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()

  

posted @ 2014-12-31 22:59  阿牧遥  阅读(151)  评论(0编辑  收藏  举报