1. 程式人生 > 實用技巧 >LeetCode173. 二叉搜尋樹迭代器

LeetCode173. 二叉搜尋樹迭代器

要實現二叉搜尋樹的迭代器,判斷當前元素是否有下一個元素和知道下一個元素是多少?
好辦,二叉搜尋樹的中序遍歷序列不是升序的嘛,那我們就可以對二叉搜尋樹中序遍歷一遍記錄到一個全域性變數數組裡面,然後要判斷是否有下一個元素和下一個元素是多少就很簡單了。

程式碼如下:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class BSTIterator {
public:
    vector<int> nodes;            //中序遍歷,把每一個節點的值記錄在nodes數組裡
    int cur;                      //當前遍歷到的是元素在nodes陣列中的下標
    int size;                     //二叉搜尋數的節點個數

    void inOrderTraversal(TreeNode* root) {            //中序遍歷二叉搜尋數,把節點的值記錄在nodes陣列中
        if(root == NULL) {
            return;
        }
        inOrderTraversal(root -> left);
        nodes.push_back(root -> val);
        inOrderTraversal(root -> right);
    }

    BSTIterator(TreeNode* root) {            //建構函式裡進行中序遍歷,初始化cur和size
        inOrderTraversal(root);          
        cur = 0;
        size = nodes.size();
    }
    
    /** @return the next smallest number */
    int next() {
        return nodes[cur++];                  //二叉搜尋數中的下一個節點值就是nodes中下一個位置的元素,要更新cur訪問下一個位置
    }
    
    /** @return whether we have a next smallest number */
    bool hasNext() {
        return cur + 1 <= size;              //只要cur不是size - 1(nodes陣列的最後一個元素),那麼他就有next
    }
};

/**
 * Your BSTIterator object will be instantiated and called as such:
 * BSTIterator* obj = new BSTIterator(root);
 * int param_1 = obj->next();
 * bool param_2 = obj->hasNext();
 */