剑指Offer 54 | 二叉搜索树的第k个结点

107 阅读2分钟

快速解数据结构,可以看看: 万字长文解说数据结构

剑指Offer系列

54.二叉搜索树的第k个结点

题目描述

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

示例1
输入

{5,3,7,2,4,6,8},3

返回值

{4}

二叉树的节点的描述如下:

public class TreeNode {
    int val = 0;
    TreeNode left = null;
    TreeNode right = null;

    public TreeNode(int val) {
        this.val = val;
    }
}

思路 & 解答

本题的思路主要是使用递归,由于二叉搜索树的每一个节点都比自己的左节点大,比自己的右节点小,那么如果求解第 k 小元素,我们不妨使用k不断减小,直到1的时候就是最小的元素。

如果root为空,那么直接返回,否则,用k减掉左子树的节点数:
如果结果为1,说明当前的root节点就是第k个节点(左子树有k-1个节点);
如果结果小于1,那么说明左子树的节点大于k-1个,第k个元素在左子树中,使用左子树进行递归;
如果结果大于1,说明左子树的节点不够k-1个,那么第k个节点肯定是在右子树中。

那么获取子树节点个数的时候,输入的是k,如果节点数大于k,结果就是负数,如果节点数小于k,结果就是正数。如果根节点为空,则直接返回k;否则先处理左子树,然后k--(表示减掉自身),然后处理右子树。

代码如下:

public class Solution62 {
    public static void main(String[] args) {
        TreeNode root = new TreeNode(5);
        root.left = new TreeNode(4);
        root.left.left = new TreeNode(3);
        root.left.left.left  = new TreeNode(2);
        root.left.left.left.left  = new TreeNode(1);
        TreeNode result = new Solution62().KthNode(root,3);
        System.out.println(result);
    }
    TreeNode KthNode(TreeNode pRoot, int k) {
        if(pRoot==null){
            return pRoot;
        }else{
            int left = getNum(pRoot.left,k);
            if(left==1){
                return pRoot;
            }else if(left<1){
                return KthNode(pRoot.left,k);
            }else{
                return KthNode(pRoot.right,left-1);
            }
        }
    }
    int getNum(TreeNode root,int k){
        if(root==null){
            return k;
        }
        int left = getNum(root.left,k);
        left--;
        return getNum(root.right,left);
    }
}

C++ 代码如下:

/**
 * struct TreeNode {
 *	int val;
 *	struct TreeNode *left;
 *	struct TreeNode *right;
 *	TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 * };
 */
class Solution {
public:
    int KthNode(TreeNode *pRoot, int k) {
        if (pRoot == NULL) {
            return -1;
        } else {
            int left = getNum(pRoot->left, k);
            if (left == 1) {
                return pRoot->val;
            } else if (left < 1) {
                return KthNode(pRoot->left, k);
            } else {
                return KthNode(pRoot->right, left - 1);
            }
        }
    }

    int getNum(TreeNode *root, int k) {
        if (root == NULL) {
            return k;
        }
        int left = getNum(root->left, k);
        left--;
        return getNum(root->right, left);
    }
};

其实根本不需要这么麻烦,因为二叉搜索树,如果使用中序遍历就可以直接得到第k个节点,也就是第k小的数值。

public class Solution {
    int count = 0;   
    int result = -1;
    public int KthNode (TreeNode proot, int k) {
        if(proot == null || k <= 0)    return -1;
        KthNode(proot.left,k);
        ++count;
        if(count == k)    return result = proot.val;
        KthNode(proot.right,k);
        return result;
    }
}

C++ 代码:

class Solution {
public:
    int count = 0;   
    int result = -1;

    int KthNode(TreeNode *proot, int k) {
        if (proot == NULL || k <= 0) {
            return -1;
        }
        KthNode(proot->left, k);
        ++count;
        if (count == k) return result = proot->val;
        KthNode(proot->right, k);
        return result;
    }
};

当然,也可以写出非递归的版本,借助堆栈即可:


public class Solution54 {
    public int KthNode(TreeNode proot, int k) {
        if (proot == null) return -1;
        Stack<TreeNode> stack = new Stack<>();
        stack.push(proot);
        TreeNode node = proot;
        int i = 0;
        while (!stack.isEmpty()) {
            //左子树
            while (node.left != null) {
                stack.push(node.left);
                node = node.left;
            }
            i++;
            if (i == k) {
                return stack.pop().val;
            }
            TreeNode temp = stack.pop();
            // 右子树
            if (temp.right != null) {
                stack.push(temp.right);
                node = temp.right;
            }
        }
        return -1;
    }
}

C++代码如下:

    int KthNode(TreeNode *proot, int k) {
        if (proot == NULL) return -1;
        stack<TreeNode *> myStack;
        myStack.push(proot);
        TreeNode *node = proot;
        int i = 0;
        while (myStack.size() > 0) {
            //左子树
            while (node->left != NULL) {
                myStack.push(node->left);
                node = node->left;
            }
            i++;
            if (i == k) {
                int val = myStack.top()->val;
                myStack.pop();
                return val;
            }
            TreeNode *temp = myStack.top();
            myStack.pop();
            // 右子树
            if (temp->right != NULL) {
                myStack.push(temp->right);
                node = temp->right;
            }
        }
        return -1;
    }