LeetCode#230 Kth Smallest Element in a BST (week7)

week7

题目

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST’s total elements.

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?
原题地址：https://leetcode.com/problems/kth-smallest-element-in-a-bst/description/

解析

题目要求找出一颗二叉搜索树中第k小的节点的值。注意到二叉搜索树中对于某个节点，其左子树所有节点的值都小于该节点的值，其右子树所有节点的值都大于该节点的值，因此我们可以采用一个栈来存储遍历过的节点，采用中序遍历的方式从根节点开始遍历，直到栈中的元素个数达到k个，此时栈顶的元素即为所求。

代码

// Definition for a binary tree node.
struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> s;
        search(root, s, k);
        return s.top()->val;
    }
    void search(TreeNode* t, stack<TreeNode*> &s, int k) {
        if (!t || s.size() >= k) {
            return;
        }
        if (t->left) {
            search(t->left, s, k);
        }
        if (s.size() < k) {
            s.push(t);
        }
        if (t->right) {
            search(t->right, s, k);
        }
    }
};