剑指Offer&LeetCode:一文详述面试中的TOP K系列问题解决方案

LeetCode 378. 有序矩阵中第K小的元素

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第k小的元素。
请注意，它是排序后的第 k 小元素，而不是第 k 个不同的元素。

示例:

matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。

提示：
你可以假设 k 的值永远是有效的, 1 ≤ k ≤ n²。

解题思路：优先队列
使用最大堆，当队列中的元素大于k时，便弹出队头元素，由于使用的是最大堆模式，所以最后返回队头元素（即为有序矩阵中第k小的元素）

• C++实现

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        priority_queue<int> pq;
        int rows = matrix.size();
        int cols = matrix[0].size();
        for(int row = 0;row < rows;row++){
            for(int col = 0; col <cols;col++){
                pq.push(matrix[row][col]);
                if(pq.size() > k){
                    pq.pop();
                }
            }
        }
        return pq.top();
    }
};

• Python

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        res = []
        rows = len(matrix)
        cols = len(matrix[0])
        for i in range(rows):
            for j in range(cols):
                res.append(matrix[i][j])
        #sorted是为了解决[[1,2],[1,3]] k=2这个测试案例
        return sorted(res)[k-1]

215. 数组中的第K个最大元素

在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例 1:

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4

说明:
你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。

解题方法：暴力法，直接升序排序之后取倒数第k个值

• C++

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        int size = nums.size();
        sort(begin(nums), end(nums));
        return nums[size - k];
        //算法的时间复杂度为 O(NlogN)，空间复杂度为 O(1)
    }
};

• Python

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        #从小到大排序，第k个最大元素即倒数第k个元素
        return sorted(nums)[-k]
        #算法的时间复杂度为 O(NlogN)，空间复杂度为 O(1)

解题思路2：快速排序，借助 partition 操作定位到最终排定以后索引为 len - k 的那个元素（特别注意：随机化切分元素

partition（切分）操作，使得：

• 对于某个索引 j，nums[j] 已经排定，即 nums[j] 经过 partition（切分）操作以后会放置在它 “最终应该放置的地方”；
• nums[left] 到 nums[j - 1] 中的所有元素都不大于 nums[j]；
• nums[j + 1] 到 nums[right] 中的所有元素都不小于 nums[j]。

partition（切分）操作总能排定一个元素，还能够知道这个元素它最终所在的位置，这样每经过一次 partition（切分）操作就能缩小搜索的范围，这样的思想叫做 “减而治之”（是 “分而治之” 思想的特例）。

切分过程可以不借助额外的数组空间，仅通过交换数组元素实现。

• Python

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        size = len(nums)

        target = size - k
        left = 0
        right = size - 1
        while True:
            index = self.partition(nums, left, right)
            if index == target:
                return nums[index]
            elif index < target:
                # 下一轮在 [index + 1, right] 里找
                left = index + 1
            else:
                right = index - 1

    #  循环不变量：[left + 1, j] < pivot
    #  (j, i) >= pivot
    def partition(self, nums, left, right):

        pivot = nums[left]
        j = left
        for i in range(left + 1, right + 1):
            if nums[i] < pivot:
                j += 1
                nums[i], nums[j] = nums[j], nums[i]

        nums[left], nums[j] = nums[j], nums[left]
        return j

最小的K个数

设计一个算法，找出数组中最小的k个数。以任意顺序返回这k个数均可。

示例：

输入： arr = [1,3,5,7,2,4,6,8], k = 4
输出： [1,2,3,4]
提示：

0 <= len(arr) <= 100000
0 <= k <= min(100000, len(arr))

解题思路：利用优先队列的思想，做法与LeetCode 378. 有序矩阵中第K小的元素类似

• C++

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        //priority_queue容器，默认从小到大排序，底层实现为最大堆
        priority_queue<int> pq;
        vector<int> res;
        for(auto x : arr){
            pq.push(x);
            if(pq.size() > k){
                pq.pop();
            }
        }
        while(!pq.empty()){
            res.push_back(pq.top());
            pq.pop();
        }
        return res;
    }
};

• Python ：利用sorted函数，一行代码解决

class Solution:
    def smallestK(self, arr: List[int], k: int) -> List[int]:
        return sorted(arr)[:k]

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

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1

3
/ \
1   4
 \
  2

输出: 1

示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3

   5
  / \
 3   6
/ \
2  4
/
1

输出: 3

解题思路：中序遍历 二叉搜索树的特点，左子树<根<右子树

• Python代码递归实现:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        res = []
        def inorder(root):
            if root:
                inorder(root.left)
                res.append(root.val)
                inorder(root.right)
        inorder(root)
        return res[k-1]

• Python使用stack实现:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        stack = []
        res = []
        while root or stack:
            # 把左子树压入栈中
            while root:
                stack.append(root)
                root = root.left
            # 输出 栈顶元素
            root = stack.pop()
            res.append(root.val)
            # 遍历右子树
            root = root.right
        return res[k-1]