【Python实战】LeetCode703、215、973: 详解 topK 系列问题

最新推荐文章于 2022-09-21 12:01:56 发布

码农谷阿莫

最新推荐文章于 2022-09-21 12:01:56 发布

阅读量417

点赞数

分类专栏： Python学习笔记文章标签： Python实现堆 leetcode heapq 算法数据结构

本文链接：https://blog.csdn.net/Mr_SCX/article/details/105345025

版权

Python学习笔记专栏收录该内容

22 篇文章 0 订阅

订阅专栏

题目描述：数据流中的第K大元素

在这里插入图片描述

解法一：排序

这个题目本身并不难解决，最容易想到的方法就是每次数据流中加入一个新元素后，就对整个数组从大到小排序，然后返回排序好的第 K 个数字即为数组中第 K 大的元素。

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.k = k
        self.nums = nums
        self.nums.sort(reverse = True) # 从大到小排序
        if len(self.nums) >= k:
            self.kLargest = self.nums[k-1] # 如果len(nums)>=k，那第 k大元素就是nums[k-1]
        else:
            self.kLargest = None # 如果len(nums) <k，那目前还不存在第 k大的元素
        
    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if val > self.kLargest: # 新元素只有大于当前数据流中的第 k大的元素才可能跻身前 k个数，否则第 k大的元素仍然不变，val永远不可能跻身前 k个数，无需加入排序数组。
            self.nums.append(val) 
            self.nums.sort(reverse = True) # 重新排序
            self.kLargest = self.nums[self.k-1] # 产生新的第 k大的元素
        return self.kLargest

可想而知，这样做的时间复杂度肯定还是比较高的，因为每插入一个大于当前第 k大元素的新数据时，就要对新的整个数组重新排序一次。

改进版本：只维护一个长度为 k的数组，每次不需要再对整个数组进行排序，而是只对前 k个数重新进行排序。

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.k = k
        self.nums = nums
        if len(nums) >= k:
            self.nums.sort(reverse = True)
            self.kLargest = self.nums[k-1]
        
    def add(self, val):
        """
        :type val: int
        :rtype: int
        """      
        if len(self.nums) < self.k: #不够 k个元素时，先插入再排序
            self.nums.append(val)
            self.nums.sort(reverse = True)
            self.kLargest = self.nums[self.k-1]
        else:
            if val > self.kLargest:
                self.nums[self.k -1] = val # 新元素大于当前数据流中的第 k大的元素时，替换掉第 k大元素为val，保持数组中只有前 k个元素
                self.nums.sort(reverse = True) # 重新对前 k个数排序
                self.kLargest = self.nums[self.k-1]
        return self.kLargest

解法二：借助小顶堆

其实，在前面的改进版本中，每次都只保留前 k个数，并只对这前 k大的数据进行维护，这其实可以借助堆这种数据结构来实现。具体来说，我们可以构建一个大小为 k 的小顶堆，保留整个数组中从大到小排序后的前 k个数，那么堆顶元素就是第 k 大的元素。每接收一个新数据，我们就拿它与当前堆顶元素比较，只要大于堆顶元素才可能加入堆，替换掉当前的堆顶元素，然后调整堆，重新找出新的第 k 大元素。

在 Python 中，内置的标准库提供了实现“二叉堆”数据结构的模块 heapq，其具体用法可以参考官方文档。heapq 其实是对一个list做原地的处理，第一个元素就是最小的，直接返回它就是小顶堆的堆顶元素。

class KthLargest(object):
    import heapq

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        # 题目中说明 len(nums) >= k-1，因此加入一个无穷小使得 len(nums) >= k，以便构造一个大小为 k 的小顶堆
        # heapq.nlargest 返回数组前 k大个数(从大到小排好序的 K个数)，返回类型为 list
        self.topK_nums = heapq.nlargest(k, nums+[float('-inf')]) 
        # 对前K大个数原地堆化，构建成一个小顶堆
        heapq.heapify(self.topK_nums) 

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        # 每次插入新数据之后，就丢弃小于第 k 大元素的数据，相当于一直维持着堆的大小为 k
        heapq.heappushpop(self.topK_nums, val)
        # 返回topK个数中的最小元素即为第K大元素，也就是这K个数构成的小顶堆的堆顶元素
        return self.topK_nums[0]

不调包，手动实现堆的写法

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        # 题目中说明 len(nums) >= k-1，因此加入一个无穷小使得 len(nums) >= k，以便构造一个大小为 k 的堆
        self.topK = sorted(nums+[float('-inf')], reverse=True)[:k] # 获取前K大个数
        self.build_heap(self.topK) # 对前K大个数原地堆化，构建成一个小顶堆
        
    def build_heap(self, nums):
        # 堆是一棵完全二叉树，对于完全二叉树来说，下标从 n/2+1 到 n 的节点都是叶子节点
        # 由于叶子节点往下堆化也只能和自己比较了
        # 从倒数第二层第一个非叶子节点起，从后往前处理数组，并且每个数据都是从上往下堆化
        for i in range(len(nums)//2, -1, -1): 
            self.heapify(nums, i)

    def heapify(self, nums, idx): # 小顶堆的堆化过程
        while True:
            min_loc = idx
            left = 2*idx + 1 # 针对下标起始为0的计算公式
            right = 2*idx + 2
            if left < len(nums) and nums[left] < nums[idx]: # 与左子节点比较
                min_loc = left
            if right < len(nums) and nums[right] < nums[min_loc]: # 与右子节点比较
                min_loc = right
            # 当前节点满足了父子节点之间的大小关系，确定了其最终所在位置，跳出循环
            if min_loc == idx: break 
            # 否则，如果不满足子节点小于等于父节点的大小关系，我们就互换两个节点
            tmp = nums[idx]
            nums[idx] = nums[min_loc]
            nums[min_loc] = tmp
            # 当发生互换后，更新min_loc为当前节点与其子节点交换后的所在位置，循环继续比较子节点与子子节点的大小关系（往下堆化）
            idx = min_loc

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if self.topK[0] < val: # 若新插入的值大于堆顶元素，则替换堆顶元素并调整堆
            self.topK[0] = val
            self.heapify(self.topK, 0)
        return self.topK[0] # topK个数构成的小顶堆的堆顶元素即为第K大元素

另一种写法

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.k = k
        self.heap = nums   
        heapq.heapify(self.heap) # heapq.heapify(x)可将一个 list 原地构建成小顶堆，这里传入的 heap其实就是个list
        while len(self.heap) > self.k: # 将堆中元素减小到只有前 k大个数
            heapq.heappop(self.heap) # 通过不断删除堆顶元素，并重新调整堆

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        if len(self.heap) < self.k:  # 堆中不够 k个元素，则直接添加 val进去
            heapq.heappush(self.heap, val) # 新增元素到堆中
        else:
            if self.heap[0] < val: # 若新的值大于堆顶元素
                heapq.heapreplace(self.heap, val) # 则删除堆顶元素，并重新堆化前 K大个数构成的小顶堆

        return self.heap[0] # 最终返回堆顶元素，即为第K大元素

数组中的第K个最大元素

在这里插入图片描述

解法一：直接从大到小排序，取第k个元素

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        return sorted(nums, reverse=True)[k-1]
        # 或：return sorted(nums)[-k]

时间复杂度：采用快排即为 O(n*logn)，n 为数组长度
空间复杂度：O(1)

解法二：维护一个大小为k的小顶堆

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        # 先取数组中前k个数构建一个小顶堆
        kLargest = nums[:k] 
        self.build_heap(kLargest) 
        # 从第k+1个数起，遍历数组中剩下的元素
        for i in range(k, len(nums)):
            if kLargest[0] < nums[i]: # 如果遍历到的元素大于当前堆顶元素，则替换堆顶元素并重新调整小顶堆
                kLargest[0] = nums[i]
                self.heapify(kLargest, 0)
        return kLargest[0] # 遍历完返回topK个数构成的小顶堆的堆顶元素即为第K大元素

    def build_heap(self, nums): # 建堆
        for i in range(len(nums)//2, -1, -1):
            self.heapify(nums, i)
        
    def heapify(self, nums, idx): # 小顶堆的堆化过程
        while True:
            min_pos = idx
            left = idx*2 + 1
            right = idx*2 + 2
            if left<len(nums) and nums[left] < nums[idx]:
                min_pos = left
            if right<len(nums) and nums[right] < nums[min_pos]:
                min_pos = right
            if min_pos == idx: break
            nums[idx], nums[min_pos] = nums[min_pos], nums[idx] # 互换
            idx = min_pos

调用Python 的 heapq 库的写法，可以精简到一行代码（果然是人生苦短，我用python，手动狗头…）

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        import heapq
        return heapq.nlargest(k,nums)[-1]

时间复杂度：往大小为 k 的堆中添加元素的时间复杂度为 O(logk)，重复该操作 n 次，故总时间复杂度为 O(n*logk)
空间复杂度：维护一个存放前 k 个元素的堆需要 O(k)

解法三：快速选择算法

本方法思路大致上与快速排序相同。简便起见，注意到第 k 个最大元素也就是第 N - k 个最小元素，因此可以用求第 N - k 小的元素的思路来解决本问题。

随机选择一个枢纽
沿着数组移动，将每个元素与枢轴进行比较，并将小于枢轴的所有元素移动到枢轴的左侧，所有大于或等于的元素都放在其右侧，最终在输出的数组中，枢轴落在了其合适的位置 pos 上。
比较 pos 和 N - k 以决定在哪边继续递归处理

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        left = 0
        right = len(nums)-1
        k_smallest = len(nums)-k # 第n-k小的元素，即为第k大元素
        return self.quick_select(nums, left, right, k_smallest)

    def partition(self, nums, left, right, pivot_index):
        pivot = nums[pivot_index]
        # 1.先把枢纽元素换到数组末尾上存放
        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

        # 2.通过游标sorted_index把nums分成两部分。
        # nums[left:sorted_index]中的元素都是小于pivot的，可以理解为已处理区间，剩下的nums[sorted_index:right]为未处理区间
        # 每次从未处理区间中取一个元素与pivot对比，如果小于pivot则将其加入到已处理区间的尾部，也就是sorted_index位置上。
        sorted_index = left
        for i in range(left, right):
            if nums[i] < pivot:
                nums[i], nums[sorted_index] = nums[sorted_index], nums[i]
                sorted_index += 1
        
        # 3.将pivot放在其最终合适的位置上
        nums[right], nums[sorted_index] = nums[sorted_index], nums[right]

        return sorted_index # 返回最终的pivot_index

    def quick_select(self, nums, left, right, target):
        if left == right: # 递归终止条件
            return nums[left]
        
        pivot_index = random.randint(left, right) # 每次随机选择一个位置作为分区的枢纽

        pivot_index = self.partition(nums, left, right, pivot_index) # 调用分区算法，确定排序后pivot_index合适的位置

        if target == pivot_index:  # 当前枢纽所在位置正好为待查找的位置
            return nums[pivot_index]
        elif target < pivot_index: # 如果待查找的位置小于当前枢纽所在位置，去左半边递归查找
            return self.quick_select(nums, left, pivot_index-1, target)
        else:                      # 如果待查找的位置大于当前枢纽所在位置，去右半边递归查找
            return self.quick_select(nums, pivot_index+1, right, target)

也可以对数组按倒序进行分区，然后直接查找第k大元素

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        left = 0
        right = len(nums)-1
        return self.quick_select(nums, left, right, k-1)

    def partition(self, nums, left, right, pivot_index):
        pivot = nums[pivot_index]
        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]

        # 每次从未处理区间中取一个元素与pivot对比，如果大于pivot则将其放到pivot_index左边。
        sorted_index = left
        for i in range(left, right):
            if nums[i] > pivot: 
                nums[i], nums[sorted_index] = nums[sorted_index], nums[i]
                sorted_index += 1

        nums[right], nums[sorted_index] = nums[sorted_index], nums[right]

        return sorted_index 

    def quick_select(self, nums, left, right, target):
        if left == right: # 递归终止条件
            return nums[left]
        
        pivot_index = random.randint(left, right) # 每次随机选择一个位置作为分区的枢纽

        pivot_index = self.partition(nums, left, right, pivot_index) # 调用分区算法，确定排序后pivot_index合适的位置

        if target == pivot_index:  # 当前枢纽所在位置正好为待查找的位置
            return nums[pivot_index]
        elif target < pivot_index: # 如果待查找的位置小于当前枢纽所在位置，去右半边递归查找
            return self.quick_select(nums, left, pivot_index-1, target)
        else:                      # 如果待查找的位置大于当前枢纽所在位置，去左半边递归查找
            return self.quick_select(nums, pivot_index+1, right, target)

时间复杂度：平均情况下为O(N)，最坏情况下为O(N²)
空间复杂度：O(1)

最接近原点的 K 个点

在这里插入图片描述

解法一：直接排序

按每个点与原点的距离从小到大排序，然后取前 k 个点即为距离原点最近的 k 个点。

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        distance_dict = {}
        for point in points:
            distance = sqrt(point[0]**2+point[1]**2)
            distance_dict[(point[0],point[1])] = distance

        sorted_dict = sorted(distance_dict.items(), key =lambda x:x[1]) # 按字典的value排序，输出类型为list

        res = []
        for i in range(K): # 取排序好的前K个点即为距离原点最近的K个点
            res.append(sorted_dict[i][0]) # sorted_dict中每一项为tuple，元组的第一个元素即为坐标点

        return res

另一种写法，可以利用python语言的匿名函数特性，通过自定义一个排序的键值比较函数来完成排序

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        sorted_points = sorted(points, key =lambda x: sqrt(x[0]**2+x[1]**2))
        return sorted_points[:K] # 前K个距离原点最近的点

时间复杂度：O(N*logN)
空间复杂度：O(1)

解法二：维护一个大小为k的大顶堆

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        import heapq
        # 取出K个距离原点最近的点
        k_smallest = heapq.nsmallest(K, points, key=lambda p: sqrt(p[0]**2+p[1]**2))
        return k_smallest

不调用库，手动实现堆排序的写法

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        # 先取数组前K个点构建一个大顶堆（根据元素与原点的距离比较大小）
        K_smallest = points[:K] 
        self.build_heap(K_smallest)
        # 从第K+1个元素开始遍历剩余的点，如果发现有比当前堆顶元素距离原点更近的点，则替换掉堆顶元素
        for i in range(K, len(points)):
            if self.calDist(K_smallest[0]) > self.calDist(points[i]):
                K_smallest[0] = points[i]
                self.heapify(K_smallest, 0)
        return K_smallest
        
    def build_heap(self, nums): # 建堆
        for i in range(len(nums)//2, -1, -1): # 从后往前遍历数组中的元素，并且每个元素都是从上往下堆化
            self.heapify(nums, i)

    def heapify(self, nums, idx): # 大顶堆的堆化过程
        while True:
            max_loc = idx
            left = idx*2+1
            right = idx*2+2
            if left<len(nums) and self.calDist(nums[left]) > self.calDist(nums[idx]):
                max_loc = left
            if right<len(nums) and self.calDist(nums[right]) > self.calDist(nums[max_loc]):
                max_loc = right
            if max_loc == idx: break
            nums[idx], nums[max_loc] = nums[max_loc], nums[idx]
            idx = max_loc

    def calDist(self, point):
        return sqrt(point[0]**2+point[1]**2)

时间复杂度：O(N*logK)
空间复杂度：O(K)

解法三：快速选择法

class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        def partition(points, left, right, pivot_idx):
            pivot = self.calDist(points[pivot_idx])
            # 先把pivot的值与最右边的值交换，存放到最右边
            points[pivot_idx], points[right] = points[right], points[pivot_idx] 
            # 把所有小于pivot的元素都放到左半边
            sorted_idx = left
            for i in range(left, right):
                if self.calDist(points[i]) < pivot:
                    points[i], points[sorted_idx] = points[sorted_idx], points[i]
                    sorted_idx += 1
            # 把pivot放到其最终合适的位置上
            points[sorted_idx], points[right] = points[right], points[sorted_idx]
            return sorted_idx

        def quick_select(points, left, right, K):
            if left == right: return # 递归终止条件，未排序区间只剩下一个元素
            # 1.随机选择一个位置作为分区pivot
            pivot_idx = random.randint(left, right) 
            # 2.调用partition函数，确定pivot最终的位置
            pivot_idx = partition(points, left, right, pivot_idx) 
            # 3.判断最小的K个值落在哪个区间中
            if K < pivot_idx-left+1: # 最小的K个值均落在pivot左侧，递归对左半边排序
                quick_select(points, left, pivot_idx-1, K)
            elif K > pivot_idx-left+1: # 已经得到（pivot_idx-left+1）个最小值，但不足K个，递归到右半边寻找剩余的K-(pivot_idx-left+1)个最小值
                quick_select(points, pivot_idx+1, right, K-(pivot_idx-left+1))
            else: # 此时pivot==K，左半边的所有值即为最小的K个值
                return

        quick_select(points, 0, len(points)-1, K) # 调用快速选择算法，找出前K个最小值
        return points[:K]

    def calDist(self, point): # 计算与原点的距离
        return sqrt(point[0]**2+point[1]**2)

时间复杂度：平均情况下为O(N)，最坏情况下为O(N²)
空间复杂度：O(1)

最后，如果大家有更好的Python解法，欢迎在评论区分享下您的解法，一起进步，感谢^ o ^~

参考：

数据结构与算法之美：堆和堆排序

https://leetcode-cn.com/problems/kth-largest-element-in-a-stream/solution/python-3xing-er-fen-otlogn-4-xing-dui-otlogk-by-qq/

https://leetcode-cn.com/problems/kth-largest-element-in-an-array/solution/shu-zu-zhong-de-di-kge-zui-da-yuan-su-by-leetcode/

码农谷阿莫

关注

0
点赞
踩
1

收藏

觉得还不错? 一键收藏
0
评论
【Python实战】LeetCode703、215、973: 详解 topK 系列问题

题目描述：数据流中的第K大元素解法一：排序这个题目本身并不难解决，最容易想到的方法就是每次数据流中加入一个新元素后，就对整个数组从大到小排序，然后返回排序好的第 K 个数字即为数组中第 K 大的元素。class KthLargest(object): def __init__(self, k, nums): """ :type k: int ...
复制链接

扫一扫

专栏目录