1. 程式人生 > >騰訊//二叉搜尋樹中第K小的元素

騰訊//二叉搜尋樹中第K小的元素

給定一個二叉搜尋樹,編寫一個函式 kthSmallest 來查詢其中第 個最小的元素。

說明:
你可以假設 k 總是有效的,1 ≤ k ≤ 二叉搜尋樹元素個數。

示例 1:

輸入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
輸出: 1

示例 2:

輸入: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
輸出: 3

進階:
如果二叉搜尋樹經常被修改(插入/刪除操作)並且你需要頻繁地查詢第 k 小的值,你將如何優化 kthSmallest

 函式?

/**
 * 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 calcTreeSize(TreeNode* root){
        if(root == NULL)
            return 0;
        return 1+calcTreeSize(root->left)+calcTreeSize(root->right);
    }
    int kthSmallest(TreeNode* root, int k) {
        if(root == NULL)
            return 0;
        int leftSize = calcTreeSize(root->left);
        if(k == leftSize + 1)
            return root->val;
        else if(leftSize >= k){
            return kthSmallest(root->left,k);
        }else{
            return kthSmallest(root->right,k-leftSize-1);
        }
    }
};