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

题目来源

• 230. 二叉搜索树中第K小的元素

题目描述

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

class Solution {

public:
    int kthSmallest(TreeNode* root, int k) {

    }
};

题目解析

morris 遍历

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int val = 0, cnt = 0;
        TreeNode *curr = root, *mostRight = nullptr;
        while (curr != nullptr){
            mostRight = curr->left;
            if(mostRight == nullptr){
                cnt++;
                cnt == k ? val = curr->val : val = val;
                curr = curr->right;
            }else{
                while (mostRight->right != nullptr && mostRight->right != curr){
                    mostRight = mostRight->right;
                }

                if(mostRight->right == nullptr){
                    mostRight->right = curr;
                    curr = curr->left;
                }else{
                    cnt++;
                    cnt == k ? val = curr->val  : val = val;
                    mostRight->right = nullptr;
                    curr =  curr->right ;
                }
            }
        }

        return val;
    }
};

分治法

在上面方法中，我们之所以需要中序遍历前k个元素，是因为我们不知道子树的节点数量。

如果我们可以记录下以每个节点为根节点的子树的节点数，并在查询第k小的值时，搜索：

• 先令node = 根节点，开始搜索

• 对当前节点node进行如下操作：

  • 如果node的左子树的节点数left < k - 1，则第k小的元素一定在node的右子树中，令node 等于右子节点，k = k - left - 1，并继续搜索
  • 如果node的左子树的节点数left = k - 1，则第k小的元素就是node，结束搜索并返回node即可
  • 如果node的左子树的节点数left > k - 1，那么第k小的元素一定在node的左子树中，令node等于node.left，并继续搜索

• 在实现中，我们既可以将以每个结点为根结点的子树的结点数存储在结点中，也可以将其记录在哈希表中。

class MyBst {
public:
    MyBst(TreeNode *root) {
        this->root = root;
        countNodeNum(root);
    }

    // 返回二叉搜索树中第k小的元素
    int kthSmallest(int k) {
        TreeNode *node = root;
        while (node != nullptr) {
            int left = getNodeNum(node->left);
            if (left < k - 1) {
                node = node->right;
                k -= left + 1;
            } else if (left == k - 1) {
                break;
            } else {
                node = node->left;
            }
        }
        return node->val;
    }

private:
    TreeNode *root;
    unordered_map<TreeNode *, int> nodeNum;

    // 统计以node为根结点的子树的结点数
    int countNodeNum(TreeNode * node) {
        if (node == nullptr) {
            return 0;
        }
        nodeNum[node] = 1 + countNodeNum(node->left) + countNodeNum(node->right);
        return nodeNum[node];
    }

    // 获取以node为根结点的子树的结点数
    int getNodeNum(TreeNode * node) {
        if (node != nullptr && nodeNum.count(node)) {
            return nodeNum[node];
        }else{
            return 0;
        }
    }
};

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        MyBst bst(root);
        return bst.kthSmallest(k);
    }
};

folloUP

类似题目

题目	思路
leetcode：面试题54. 二叉搜索树的第k大节点
leetcode：230. 二叉搜索树中第K小的元素 Kth Smallest Element in a BST
leetcode：94. 二叉树的中序遍历 Binary Tree Inorder Traversal	Morris遍历
leetcode：671. 二叉树中第二小的节点 Second Minimum Node In a Binary Tree