Code & Func

貌似又是一道之前做了，但是没写题解的题目。

这道题要求我们按从小到大的顺序返回二叉搜索树的值，而我们知道二叉搜索树的中序遍历就是从小到大的，所以问题就变成了，对一个二叉树的中序遍历问题。

因为之前总结过二叉树遍历，所以这里我们可以套用当时提到的三种方法来解这道题：

递归

由于题目只要求了next()和hasNext()的时间复杂度，所以我们可以在构造器中对二叉树进行遍历，然后存储下来：

1
class BSTIterator {
2
public:
3
    vector<int> vec;
4
    int index;
5
    BSTIterator(TreeNode* root) {
6
        index = 0;
7
        inorderTraversal(root);
8
    }
9

10
    void inorderTraversal(TreeNode *root) {
11
        if (root) {
12
            inorderTraversal(root->left);
13
            vec.push_back(root->val);
14
            inorderTraversal(root->right);
15
        }
12 collapsed lines
16
    }
17

18
    /** @return the next smallest number */
19
    int next() {
20
        return vec[index++];
21
    }
22

23
    /** @return whether we have a next smallest number */
24
    bool hasNext() {
25
        return index != vec.size();
26
    }
27
};

这样next()和hasNext()的时间复杂度肯定是O(1)的，但是空间复杂度却是O(n)，而题目要求的是O(h)。

基于栈进行迭代。

基于栈对二叉树进行迭代的中序遍历大体可以分为两部：

不断地把先左节点移动，并把节点压入栈中（以下简称 step 1)。
弹出栈顶节点，并输出，然后先右节点方向移动（以下简称 step 2)。

一旦做完以上两步，我们就输出了一个值。

因为当时是在一个循环中实现的，所以和现在的情况是不一样的，root是在构造器中输入的，所以我们得在构造器中就把root压入栈中，为了不把过程弄的复杂，所以我们在构造器中就直接把 step 1 给做完，然后在next()中把两步的顺序倒过来，即先执行 step 2 然后执行 step 1，当然 step 2 中的输出操作显然要放到最后来做。因为在两个地方都进行了 step 1，所以我们可以把它抽象成一个单独的函数pushleft：

1
class BSTIterator {
2
public:
3
    stack<TreeNode *> st;
4
    BSTIterator(TreeNode* root) {
5
        pushleft(root);
6
    }
7

8
    void pushleft(TreeNode *root) {
9
        while(root) {
10
            st.push(root);
11
            root = root->left;
12
        }
13
    }
14

15
    /** @return the next smallest number */
11 collapsed lines
16
    int next() {
17
        auto root = st.top(); st.pop();
18
        pushleft(root->right);
19
        return root->val;
20
    }
21

22
    /** @return whether we have a next smallest number */
23
    bool hasNext() {
24
        return !st.empty();
25
    }
26
};

因为基于栈的遍历算法的空间复杂度是O(h)，所以上面这个算法的空间复杂度也是O(h)（其实也很好理解，因为栈要临时存放节点个数最大就是 h），然后hasNext()的时间复杂度显然也是O(1)。下面的问题就是，next()的时间复杂度是否是O(1)。

如果我们仔细观察一下题目的话，我们会发现它要求的是平均时间复杂度，因为一颗有n个节点的树进行遍历，我们需要做n次st.push，同时next()我们也要调用n次才能遍历整棵树，所以n / n = 1，即平均时间复杂度为O(1)。

莫里斯遍历

我们先看下莫里斯遍历的代码：

1
TreeNode *GetRightLeaf(TreeNode *root, TreeNode *end) {
2
    while(root->right && root->right != end) root = root->right;
3
    return root;
4
}
5

6
void inorderTraversal(TreeNode* root) {
7
    while(root) {
8
        if (root->left) {
9
            TreeNode *p = GetRightLeaf(root->left, root);
10
            if (p->right == root) {
11
                p->right = nullptr;
12
                cout << root->val << endl;
13
                root = root->right;
14
            } else {
15
                p->right = root;
8 collapsed lines
16
                root = root->left;
17
            }
18
        } else {
19
            cout << root->val << endl;
20
            root = root->right;
21
        }
22
    }
23
}

几乎可以什么都不用改的情况下，把代码移植过去：

1
class BSTIterator {
2
public:
3
    TreeNode *root;
4
    TreeNode *GetRightLeaf(TreeNode *root, TreeNode *end) {
5
        while(root->right && root->right != end) root = root->right;
6
        return root;
7
    }
8

9
    BSTIterator(TreeNode* _root):root(_root) {
10

11
    }
12

13
    /** @return the next smallest number */
14
    int next() {
15
        int res = -1;
28 collapsed lines
16
        while(root) {
17
            if (root->left) {
18
                TreeNode *p = GetRightLeaf(root->left, root);
19
                if (p->right == root) {
20
                    p->right = nullptr;
21
                    // cout << root->val << endl;
22
                    res = root->val;
23
                    root = root->right;
24
                    break;
25
                } else {
26
                    p->right = root;
27
                    root = root->left;
28
                }
29
            } else {
30
                // cout << root->val << endl;
31
                res = root->val;
32
                root = root->right;
33
                break;
34
            }
35
        }
36
        return res;
37
    }
38

39
    /** @return whether we have a next smallest number */
40
    bool hasNext() {
41
        return root != nullptr;
42
    }
43
};

这个的空间复杂度显然是O(1)，而时间复杂度，我们可以这样理解，每个节点都会被访问两遍，即O(2n)，而next()要调用 n 次，所以时间复杂度是O(2n / 2) = O(2) = O(1)。