[数据结构与算法]【Leetcode】230. 二叉搜索树中第K小的元素

题目描述

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

题解

大顶堆法

（前序）遍历树的所有元素，设定一个大顶堆maxHeap，如果堆的size没达到k，就一直往里存元素，如果size达到了k，则凡是小于堆顶元素才能被存入maxHeap。

执行用时：1 ms, 在所有 Java 提交中击败了43.60%的用户

内存消耗：38.2 MB, 在所有 Java 提交中击败了59.98%的用户

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    PriorityQueue<Integer> maxHeap;
    int k;
        
    public int kthSmallest(TreeNode root, int k) {
        this.maxHeap = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer a, Integer b) {
                return b.compareTo(a);
            }
        });
        this.k = k;
        
        preOrder(root);
        return maxHeap.peek();
    }
    
    public void preOrder(TreeNode root) {
        if (root == null)
            return;
        if (maxHeap.size() >= k) {
            if (maxHeap.peek() > root.val) {
                maxHeap.poll();
                maxHeap.add(root.val);
            }
        }
        else {
            maxHeap.add(root.val);
        }
        preOrder(root.left);
        preOrder(root.right);
    }
}

中序遍历法

BST的中序遍历得到的序列就是升序的，那就直接数k个数，数到第k个返回就行了。

执行用时：1 ms, 在所有 Java 提交中击败了43.60%的用户

内存消耗：38 MB, 在所有 Java 提交中击败了89.59%的用户

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    int count = 0;
    int res = 0;
        
    public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        TreeNode node = root;
        
        while (node != null || !stack.isEmpty()) {
            while (node != null) {
                stack.push(node);
                node = node.left;
            }
            if (!stack.isEmpty()) {
                node = stack.pop();
                count++;
                if (count == k) {
                    res = node.val;
                    break;
                }
                node = node.right;
            }
        }
        
        return res;
    }
}