LeetCode 230：二叉搜索树中第K小的元素

 最基础的解法：将整个二叉树中序遍历存入序列，取第k个(index = k-1)返回。如果是第k大就反向中序遍历。 

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if(!root) return 0;
        InorderTraversal(root);
        return seq[k-1];
    }
private:
    vector<int> seq;
    void InorderTraversal(TreeNode* node) {
        if(!node) return;
        if(node->left) {
            InorderTraversal(node->left);
        }
        seq.push_back(node->val);
        if(node->right) {
            InorderTraversal(node->right);
        }
    }
};

改进：在遍历到第k个时就结束遍历，返回值。用迭代的方式遍历：

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if(!root) return 0;
        stack<TreeNode*> s;
        TreeNode* temp;
        while(root || !s.empty()) {
            if(root) {
                s.push(root);
                root = root->left;
            }
            else {
                temp = s.top();
                s.pop();
                if(--k == 0) break;
                root = temp->right;
            }
        }
        return temp->val;
    }
};

改进 2：利用递归，巧妙利用引用来实现提前终止递归

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if(!root) return 0;

        int res;
        int flag = 0;
        dfs(root, res, k, flag);
        return res;
    }
private:
    void dfs(TreeNode* root, int& res, int& k, int& flag) {
        if(!root || flag) return;
        //if(root->left) dfs(root->left, res, --k, flag); //我之前程序就是错在这里，在递归调用时k--是不对的，访问到最底层开始往回返的时候再让k--
        if(root->left) dfs(root->left, res, k, flag);
        if(--k == 0) {
            res = root->val;
            flag = 1;
            return;
        }
        if(root->right) dfs(root->right, res, k, flag);
    }
};

利用Python中的generator来实现递归的提前终止

class Solution:
    def mid_order(self, root: TreeNode) -> int:
        if not root: return
        yield from self.mid_order(root.left)
        yield root.val
        yield from self.mid_order(root.right)
    
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        gene = self.mid_order(root)
        for _ in range(k-1):
            next(gene)
        return next(gene)