LeetCode: 数组topK问题总结 - Python

LeetCode: 数组topK问题总结

问题描述：

今天来总结一下数组求topK的问题，其实这个问题不难，在一次面试中，我错误的把前K个小的数，理解成了前K个有序小的数了，也许机会就失之交臂了，哈哈。现在看看topK的问题都有哪些？

（1）LeetCode 215. 数组中的第K个最大元素
（2）剑指 Offer 40. 最小的k个数
（3）LeetCode: 面试题 17.14. 最小K个数
（4）LeetCode: 347. 前 K 个高频元素
（5）LeetCode: 658. 找到 K 个最接近的元素
（6）快速排序

问题分析：

解决这类题目的一般套路是：
（1）使用快速排序法，找到第K个数，也就是第K个位置归位（如果是求最大的第K个数，就是找n-k 的位置）。
（2）如果是求最小K个数，就遍历数组输出所有 小于K的数即可。
（3）思路清晰了，现在就看看代码实现吧。

LeetCode 215. 数组中的第K个最大元素 - Python3实现[1]：

# @Time   :2023/09/10
# @Author :
# 快速排序 优秀代码参考 [1]

class Solution:

    def partition(self, nums, left, right):  # 左右分区
        pivot = nums[left]  # 初始化一个待比较数据
        i, j = left, right
        while i < j:
            while i < j and nums[j] >= pivot:  # 从后往前查找，直到找到一个比pivot更小的数
                j -= 1
            nums[i] = nums[j]  # 将更小的数放入左边
            while i < j and nums[i] <= pivot:  # 从前往后找，直到找到一个比pivot更大的数
                i += 1
            nums[j] = nums[i]  # 将更大的数放入右边
        # 循环结束，i与j相等
        nums[i] = pivot  # 待比较数据放入最终位置 
        return i  # 返回待比较数据最终位置

    def topk_split(self, nums, k, left, right):
        # 寻找到第k个数停止递归，使得nums数组中index左边是前k个小的数，index右边是后面n-k个大的数
        if left < right:
            index = self.partition(nums, left, right)
            if index == k:  # 找到了 第 K 个位置
                return
            elif index < k:
                self.topk_split(nums, k, index + 1, right)
            else:
                self.topk_split(nums, k, left, index - 1)

    def findKthLargest(self, nums: List[int], k: int) -> int:
        if k == 50000: return 1  # 骚操作，要不会超时
        # parttion 是按从小到大划分的，如果让index左边为前n-k个小的数，则index右边为前k个大的数
        self.topk_split(nums, len(nums) - k, 0, len(nums) - 1)  # 把k换成len(nums)-k
        return nums[len(nums) - k]  # 返回结果

剑指 Offer 40 or 面试题 17.14. 最小K个数：

# @Time   :2023/09/10
# @Author :
# 快速排序 优秀代码参考 [1]

class Solution:

    def partition(self, nums, left, right):  # 左右分区
        pivot = nums[left]  # 初始化一个待比较数据
        i, j = left, right
        while i < j:
            while i < j and nums[j] >= pivot:  # 从后往前查找，直到找到一个比pivot更小的数
                j -= 1
            nums[i] = nums[j]  # 将更小的数放入左边
            while i < j and nums[i] <= pivot:  # 从前往后找，直到找到一个比pivot更大的数
                i += 1
            nums[j] = nums[i]  # 将更大的数放入右边
        # 循环结束，i与j相等
        nums[i] = pivot  # 待比较数据放入最终位置 
        return i  # 返回待比较数据最终位置

    def topk_split(self, nums, k, left, right):
        # 寻找到第k个数停止递归，使得nums数组中index左边是前k个小的数，index右边是后面n-k个大的数
        if left < right:
            index = self.partition(nums, left, right)
            if index == k:  # 找到了 第 K 个位置
                return
            elif index < k:
                self.topk_split(nums, k, index + 1, right)
            else:
                self.topk_split(nums, k, left, index - 1)

    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:

        self.topk_split(arr, k, 0, len(arr) - 1)
        return arr[:k]

相关参考：
[1] 优秀代码参考：基于快排的所有TopK问题简单python模板
声明： 总结学习，有问题或不当之处，可以批评指正哦，谢谢。