剑指54：二叉搜索树的第k大节点

传送门

nowcoder
leetcode

题目描述

给定一棵二叉搜索树，请找出其中的第 k 小的结点。

C++ 代码 - nowcoder

/*
  	中序遍历，从小到大的排列顺序
    迭代写法
*/
class Solution {
public:
    int KthNode(TreeNode* pRoot, int k) {
        if (pRoot == nullptr) return -1;

        stack<TreeNode*> stk;
        while (!stk.empty() || pRoot != nullptr) {
            while (pRoot != nullptr) {
                stk.push(pRoot);
                pRoot = pRoot->left;
            }
            pRoot = stk.top(); stk.pop();
            if (-- k == 0) return pRoot->val;
            pRoot = pRoot->right;
        }

        return -1;
    }
};

C++ 代码 - leetcode

/*
	本题可以根据 BST 的中序遍历计算第 k 大的元素。
	常规的中序遍历得到的顺序是从小到大的，而我们需要从大到小的顺序。
	把中序遍历框架反过来，先递归遍历右子树，再递归遍历左子树，即可获得降序结果。
*/
class Solution {
private:
    int res = 0;  // 记录结果
    int rank = 0; // 记录当前元素的排名

public:
    int findTargetNode(TreeNode* root, int k) {
        // 利用 BST 的中序遍历特性
        traverse(root, k);
        return res;
    }

    void traverse(TreeNode* root, int k) {
        if (root == nullptr) {
            return;
        }
        traverse(root->right, k);

        /* 中序遍历代码位置 */
        rank ++;
        if (k == rank) {
            // 找到第 k 大的元素
            res = root->val;
            return;
        }
        /*****************/
        
        traverse(root->left, k);
    }
};