1. 程式人生 > >173 Binary Search Tree Iterator

173 Binary Search Tree Iterator

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.  呼叫next()返回二叉查詢樹中的最小的元素。

Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.  注意:next()和hasNext()執行的平均時間

複雜度為O(1),空間複雜度為O(h),其中h是樹的高度。

解題思路 如果不看注意的話,一個非常簡單的解決辦法是遍歷這棵二叉樹,把二叉樹的元素以從大到小的方式放到一個棧裡,這樣next()從棧頂獲取元素,hasNext()呼叫棧不為空的判定方法。這種方法能滿足題目中的時間複雜度要求,但是O(n)的空間複雜度無法滿足題目中O(h)空間複雜度要求。

我們維護一個棧,暫且稱之為最小元素棧,棧中存放的是結點指標,向棧中新增元素髮生在兩種情況下:1.初始化迭代器時;2.呼叫next()從棧中彈出一個元素時。我們結合下面的例項來描述這個棧中元素的新增和刪除過程。對於hasNext()方法,呼叫棧的empty()方法即可,不做具體分析,僅分析next()方法以及迭代器的資料結構。

綜上所述,我們在棧中儲存的是一部分極小值,且棧頂一定是當前樹中未訪問過的結點中的最小結點。在呼叫next()彈出最小結點時,需要把被彈出結點的右子樹新增到棧中,也就是把比當前最小結點大但是比次小結點小的那些結點中的一部分(不是全部,只是左孩子分支上的一部分)新增到棧中。

然後我們來看,我們的這個資料結構是否滿足題目要求。我們可以看到每一個結點一定要進入棧中,且該操作是在構造過程和next()中執行的,其中構造過程只執行一次,而next()執行多次。但是進棧的操作存在差異性,有時候呼叫next()不會往棧中新增任何元素,有時候會新增多個,但是至多不超過h個,其中h是樹的高度,因為新增時只順著樹的左分支往下新增。我們一共要呼叫n次next(),在這n次呼叫中,每個結點都進棧/出棧,共n次入棧操作,n次出棧操作,平均起來,每次next()只有一次入棧操作+一次出棧操作,因此滿足平均時間複雜度O(1)的要求。再看空間複雜度,剛才我們已經說過,在往棧中新增元素時只順著樹的左分支往下新增,因此至多新增h個結點,其中h就是樹的高度,因此滿足空間複雜度O(h)的要求。

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */

public class BSTIterator {
    
    Stack<TreeNode> stack;
    
    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        if(root == null)
            return;
        stack.push(root);
        TreeNode cur = root.left;
        while(cur != null){
            stack.push(cur);
            cur = cur.left;
        }
    }
    
    

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !stack.isEmpty();
    }

    /** @return the next smallest number */
    public int next() {
        
        TreeNode res = stack.pop();
        TreeNode cur = res.right;
        while(cur != null){
            stack.push(cur);
            cur = cur.left;
        }
        return res.val;
    }
}

/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = new BSTIterator(root);
 * while (i.hasNext()) v[f()] = i.next();
 */