[Leetcode] Binary Search Tree Iterator 二叉搜素樹迭代器

shixinzhang 發(fā)布于2019-08-14 12:36 / 2783人閱讀

摘要：代碼先找到第一個節(jié)點，并把路徑入棧棧為空時不再有下一個棧頂是待返回元素如果該元素有右節(jié)點，把它的右節(jié)點及右節(jié)點的所有左邊節(jié)點都壓入棧中

Binary Search Tree Iterator 最新更新：https://yanjia.me/zh/2019/02/...

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.Calling next() will return the next smallest number in the BST.
Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.

棧法 復雜度

時間 O(N) 空間 O(1)

思路

這道題和Inorder Successor in BST很像，區(qū)別在于Successor那題是給定節(jié)點求下一個，而這題是要做一個迭代器，返回當前指向的值，其實就是求上次返回的節(jié)點的下一個，本質(zhì)是一樣的。題目要求最多只能用O(H)的空間，所以思路也和那題一樣，用一個Stack記錄從根節(jié)點到當前節(jié)點的路徑。next的時候就返回Stack最上面的元素。不過拿出最上面的元素后，我們還要看一下這個被返回的元素是否有右節(jié)點，如果有的話，就把它的右節(jié)點及右節(jié)點的所有左邊節(jié)點都壓入棧中。另外，初始化棧時，我們要找到最左邊的節(jié)點，也就是中序遍歷的第一個節(jié)點，并把根到第一個節(jié)點的路徑都壓入棧。

這題還可以結(jié)合Binary Tree Inorder Traverse的迭代做法一起理解。他們的共同點都是用一個棧，將最左邊的節(jié)點都壓入棧。

代碼

public class BSTIterator {
    
    Stack stk;

    public BSTIterator(TreeNode root) {
        stk = new Stack();
        // 先找到第一個節(jié)點，并把路徑入棧
        while(root != null){
            stk.push(root);
            root = root.left;
        }
    }

    public boolean hasNext() {
        // 棧為空時不再有下一個
        return !stk.isEmpty();
    }

    public int next() {
        // 棧頂是待返回元素
        TreeNode curr = stk.pop();
        int res = curr.val;
        // 如果該元素有右節(jié)點，把它的右節(jié)點及右節(jié)點的所有左邊節(jié)點都壓入棧中
        if(curr.right != null){
            curr = curr.right;
            while(curr != null){
                stk.push(curr);
                curr = curr.left;
            }
        }
        return res;
    }
}