天天看點

[每日一題] 134. 二叉搜尋樹中第K小的元素(BST樹、遞歸、搜尋技巧、多方法)

文章目錄

    • 1. 題目來源
    • 2. 題目說明
    • 3. 題目解析
          • 方法一:BST性質、非遞歸中序周遊
          • 方法二:遞歸中序周遊
          • 方法三:分治法
          • 方法四:統計左右子樹節點個數、探索二叉搜尋樹解法(絕妙)

1. 題目來源

連結:二叉搜尋樹中第K小的元素

來源:LeetCode

2. 題目說明

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

說明:

你可以假設 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

函數?

3. 題目解析

方法一:BST性質、非遞歸中序周遊

用 BST 的性質來解題,最重要的性質是就是左<根<右,如果用中序周遊所有的節點就會得到一個有序數組。是以解題的關鍵還是中序周遊啊。

先來看一種非遞歸的方法,中序周遊最先周遊到的是最小的結點,隻要用一個計數器,每周遊一個結點,計數器自增1,當計數器到達 k 時,傳回目前結點值即可。

參見代碼如下:

// 執行用時 :24 ms, 在所有 C++ 送出中擊敗了70.31%的使用者
// 記憶體消耗 :21.6 MB, 在所有 C++ 送出中擊敗了33.58%的使用者

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int cnt = 0;
        stack<TreeNode*> s;
        TreeNode *p = root;
        while (p || !s.empty()) {
            while (p) {
                s.push(p);
                p = p->left;
            }
            p = s.top(); 
            s.pop();
            ++cnt;
            if (cnt == k) 
                return p->val;
            p = p->right;
        }
        return 0;
    }
};
           
方法二:遞歸中序周遊

當然,此題也可以用遞歸來解,還是利用中序周遊來解。

參見代碼如下:

// 執行用時 :28 ms, 在所有 C++ 送出中擊敗了44.03%的使用者
// 記憶體消耗 :22 MB, 在所有 C++ 送出中擊敗了6.89%的使用者

/**
 * 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 InOrder(TreeNode* root, int& k) {
        if (root) {
            int val = InOrder(root->left, k);
            if (k == 0) 
                return val;
            if (--k == 0) 
                return root->val;
            return InOrder(root->right, k);
        }
        return -1;
    }

    int kthSmallest(TreeNode* root, int k) {
        return InOrder(root, k);
    }
};
           
方法三:分治法

來看一種分治法的思路,由于 BST 的性質,可以快速定位出第 k 小的元素是在左子樹還是右子樹,即有:

  • 首先計算出左子樹的結點個數總和 cnt
  • 如果 k 小于等于左子樹結點總和 cnt,說明第 k 小的元素在左子樹中,直接對左子結點調用遞歸即可
  • 如果 k 大于 cnt + 1,說明目标值在右子樹中,對右子結點調用遞歸函數

注意此時的 k 應為 k - cnt - 1,因為已經減少了 cnt + 1 個結點。如果 k 正好等于 cnt + 1,說明目前結點即為所求,傳回目前結點值即可。

參見代碼如下:

// 執行用時 :20 ms, 在所有 C++ 送出中擊敗了88.64%的使用者
// 記憶體消耗 :22.1 MB, 在所有 C++ 送出中擊敗了5.23%的使用者

/**
 * 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) {
        int cnt = count(root->left);
        if (k <= cnt) 
            return kthSmallest(root->left, k);
        else if (k > cnt + 1) 
            return kthSmallest(root->right, k - cnt - 1);
        return root->val;
    }
    int count(TreeNode* node) {
        if (!node) 
            return 0;
        return 1 + count(node->left) + count(node->right);
    }
};
           
方法四:統計左右子樹節點個數、探索二叉搜尋樹解法(絕妙)

這道題的 進階 中說假設該 BST 被修改的很頻繁,而且查找第k小元素的操作也很頻繁,問我們如何優化。

其實最好的方法還是像上面的解法那樣利用分治法來快速定位目标所在的位置,但是每個遞歸都周遊左子樹所有結點來計算個數的操作并不高效,是以應該修改原樹結點的結構,使其儲存包括目前結點和其左右子樹所有結點的個數,這樣就可以快速得到任何左子樹結點總數來快速定位目标值了。

定義了新結點結構體,然後就要生成新樹,還是用遞歸的方法生成新樹,注意生成的結點的 count 值要累加其左右子結點的 count 值。

然後在求第 k 小元素的函數中,先生成新的樹,然後調用遞歸函數。在遞歸函數中,不能直接通路左子結點的 count 值,因為左子節結點不一定存在,是以要先判斷,如果左子結點存在的話,那麼跟上面解法的操作相同。如果不存在的話,當此時 k 為 1 的時候,直接傳回目前結點值,否則就對右子結點調用遞歸函數,k自減1。

在 LeetCode探索 二叉搜尋樹 對該種搜尋方法進行了诠釋,真的絕妙的思想。

[每日一題] 134. 二叉搜尋樹中第K小的元素(BST樹、遞歸、搜尋技巧、多方法)
[每日一題] 134. 二叉搜尋樹中第K小的元素(BST樹、遞歸、搜尋技巧、多方法)

參見代碼如下:

// 執行用時 :32 ms, 在所有 C++ 送出中擊敗了24.55%的使用者
// 記憶體消耗 :24.9 MB, 在所有 C++ 送出中擊敗了5.23%的使用者

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */


/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */


 // Follow up
class Solution {
public:
    struct MyTreeNode {
        int val;
        int count;
        MyTreeNode *left;
        MyTreeNode *right;
        MyTreeNode(int x) : val(x), count(1), left(NULL), right(NULL) {}
    };
    
    MyTreeNode* build(TreeNode* root) {
        if (!root) 
            return NULL;
        MyTreeNode* node = new MyTreeNode(root->val);
        node->left = build(root->left);
        node->right = build(root->right);
        if (node->left) 
            node->count += node->left->count;
        if (node->right) 
            node->count += node->right->count;
        return node;
    }
    
    int kthSmallest(TreeNode* root, int k) {
        MyTreeNode *node = build(root);
        return helper(node, k);
    }
    
    int helper(MyTreeNode*& node, int k) {
        if (node == nullptr)
            return 0;
        if (node->left) {
            int cnt = node->left->count;
            if (k <= cnt) 
                return helper(node->left, k);
            else if (k > cnt + 1) 
                return helper(node->right, k - 1 - cnt);
            return node->val;
        } 
        else {
            if (k == 1) 
                return node->val;
            return helper(node->right, k - 1);
        }
    }
};