【剑指offer】54、二叉搜索树的第k大的节点

题目

给定一颗二叉搜索树，找出其中第k大的节点（从小到大），返回其节点

思路一

最朴素的想法，中序遍历就是二叉搜索树的递增序列，那直接写出整棵树的中序遍历即可

递归版本

class Solution {
public:
    TreeNode* KthNode(TreeNode* pRoot, int k)
    {
        if (!pRoot || k<=0)
            return nullptr;
        vector<TreeNode*> in_order;
        InOrder(pRoot, in_order);
        if (k > in_order.size())  // 不加这句话会下标溢出
            return nullptr;
        return in_order[k-1];
    }

    void InOrder(TreeNode* pRoot, vector<TreeNode*>& in_order){
        if (pRoot == nullptr)
            return;
        InOrder(pRoot->left, in_order);
        in_order.push_back(pRoot);
        InOrder(pRoot->right, in_order);
    }
};

同时也应该掌握非递归版本

class Solution {
public:
    TreeNode* KthNode(TreeNode* pRoot, int k)
    {
        if (!pRoot || k<=0)
            return nullptr;
        vector<TreeNode*> in_order;
        stack<TreeNode*> node;
        TreeNode* temp = pRoot;
        while (!node.empty() || temp)
        {
            if(temp)
            {
                node.push(temp);
                temp = temp->left;
            }
            else{
                temp = node.top();
                in_order.push_back(temp);
                node.pop();
                temp = temp->right;
            }
        }
        if (k > in_order.size())
            return nullptr;
        return in_order[k-1];
    }
};

思路二

我们只需得到从小到大的第k个数，没有必要全部遍历。根据中序的思想，可以用一个count计数。

递归版本

class Solution {
public:
    int count = 0;
    TreeNode* KthNode(TreeNode* pRoot, unsigned int k)
    {
        if(pRoot){ 
            TreeNode* res = KthNode(pRoot->left, k);
            if(res)
                return res;
            count++;
            if(count == k)
                return pRoot;
            res = KthNode(pRoot->right,k);
            if(res)
                return res;
        }
        return nullptr;
    }
};

同样只需修改非递归版本的中序遍历即可，每次压入vec时相当于找到一个节点，此时count++即可。

class Solution {
public:
    TreeNode* KthNode(TreeNode* pRoot, int k)
    {
        if (!pRoot || k<=0)
            return nullptr;
        int count = 0;
        stack<TreeNode*> node;
        TreeNode* temp = pRoot;
        while (!node.empty() || temp)
        {
            if(temp)
            {
                node.push(temp);
                temp = temp->left;
            }
            else{
                temp = node.top();
                count++;
                if (count == k)   //修改这部分就可以，而且非常好理解
                    return temp;
                node.pop();
                temp = temp->right;
            }
        }
        return nullptr;
    }
};