第K大的数

· 约 2 分钟阅读 · 次阅读 leetcode

快排划分的思想。

class Solution:
    def partition(self, nums, left, right):
        pivot = left
        index = left + 1
        for i in range(index, right+1):
            # print('left,right,i', left, right, i)
            if nums[i] > nums[pivot]:
                self.swap(nums, i, index)
                index += 1
        self.swap(nums, pivot, index-1)
        return index-1
            
    def swap(self, nums, i, j):
        nums[i],nums[j] = nums[j], nums[i]
    
    def partitionK(self, nums, left, right, k):
        index = self.partition(nums, left, right)
        length = index-left+1
        # 如果已划分的个数length大于需要找的数k, 则在index左半部分重新找k个数
        if length > k:
            self.partitionK(nums, left, index-1, k)
        # 如果已划分的个数length小于需要找的数k, 则在index右半部分再找k-length个数即可
        if length < k:
            self.partitionK(nums, index+1, right, k-length)

    def findKthLargest(self, nums: List[int], k: int) -> int:
        self.partitionK(nums, 0, len(nums)-1, k)
        return nums[k-1]

上面的partitionK可以改为迭代的写法

def partitionK(self, nums, left, right, k):
        while left < right:
            index = self.partition(nums, left, right)
            length = index-left+1
            if k == length:
                break
            if k < length:
                right = index -1
            if k > length:
                left = index + 1
                k = k - length

堆排序的思想

class Solution:
    # 构建小顶堆
    def adjust(self, nums, i, length):
        k = 2*i + 1
        tmp = nums[i]
        while k < length:
            if k + 1 < length and nums[k+1] < nums[k]:
                k += 1
            if k < length and nums[k] < tmp:
                nums[i] = nums[k]
                i = k
            k = 2*k + 1
        nums[i] = tmp

    def buildHeap(self, nums):
        for i in range(len(nums)//2-1, -1, -1):
            self.adjust(nums, i, len(nums)) 
        return nums 

    def swap(self, nums, i, j):
        nums[i],nums[j] = nums[j], nums[i]
    
    # 构建小顶堆,堆中存储最大的前k个数,堆顶就是第K大的数
    # 然后循环遍历剩下的数,如果比堆顶元素大,则将之与堆顶元素交换并调整
    # 返回堆顶元素
    def findKthLargest(self, nums: List[int], k: int) -> int:
        nums[:k] = self.buildHeap(nums[:k])
        for i in range(k, len(nums)):
            if nums[i] > nums[0]:
                self.swap(nums, i, 0)
                self.adjust(nums, 0, k)
        return nums[0]

当然,也可以完全参考堆排序的过程,不过只输出前k个数