剑指 Offer II 076. 数组中的第 k 大的数字

Posted 2021-10-05 炫云云

剑指 Offer II 076. 数组中的第 k 大的数字

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

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

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

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

堆排序

07堆排序 python

总体思路：维护一个大小为$k$的最小堆，堆顶是这$k$个数里的最小的，从k开始遍历数组，每次删除堆顶数，剩余k个数，后返回堆顶元素即可。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = []
        for num in nums:
            heapq.heappush(heap, num)
            if len(heap) > k:
                heapq.heappop(heap)
        return heap[0]

# 自己实现堆
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        heap = nums[:k]
        # 建立含k个元素的小根堆
        for i in range((k-2)//2, -1, -1):
            self.heapify(heap , i , k-1)
        #若k之后的元素大于根节点，则将该元素与根节点替换，然后做一次调整
        for j in range(k,n):
            if nums[j] > heap[0]: #找前k大的数
                heap[0] = nums[j]
                self.heapify(heap, 0 , k-1)
        return heap[0] #堆顶就是第k大的数了

    # 建立小根堆
    def heapify(self,nums, index, end):
        left = index * 2 + 1
        right = left + 1
        while left <= end:
            # 当前节点为非叶子结点
            min_index = index
            if nums[left] < nums[min_index]:
                min_index = left
            if right <= end and nums[right] < nums[min_index]:
                min_index = right
            if index == min_index:
                # 如果不用交换，则说明已经交换结束
                break
            # 若子树的值比较小，则根节点换成子树，然后向下看一层
            nums[index], nums[min_index] = nums[min_index], nums[index]
            # 继续调整子树
            index = min_index
            left = index * 2 + 1
            right = left + 1

手写堆

堆的本质是一个完全二叉树，新加入的元素先加到末尾，然后按规则上浮；弹出元素后，堆顶元素赋值为末尾元素，然后按规则下沉；最核心的是上浮和下沉两个方法，之前一直畏难没有去手写，这次总算干掉了它，其实并不复杂

findKthLargest方法：

和前面直接调API的方法差不多

1. 先把当前数push进去，如果push失败，弹出一个元素，再push。因为我们维护的是一个大小为$k+1$的堆，和维护大小为k的堆相比，可以不用比较加入元素和堆顶元素的大小，只要在最后返回前令堆大小为$k$即可。
2. 如果堆的大小为$k+1$​，pop出堆顶元素
3. 返回这个大小为k的堆中的堆顶元素

手写堆Heap成员：

• 成员变量：

• • heap：一个长度为k + 2的数组，解释一下为什么长度为k + 2：

  • • 首先，为了索引计算的方便，索引为0的位置不存储元素，而是从索引为1开始存储；可以自己画一个树试试，如果从索引1开始的话，对于索引为 i 的节点，有如下规律：

      父节点索引为 i >> 1

      左子节点索引为 i << 1

      右子节点索引为 i << 1 | 1

    • 其次，为了写代码方便，我们维护的是一个大小为$k+1$的堆，而不是$k$；也就是堆中其实多一个元素，在要返回之前先pop出一个元素，让堆的大小为$k$​了再返回。

    • 综上，总共要多两个位置，因此长度为$k+2$

  • size：记录当前heap已存储了多少元素

• 方法：

• • push：加入一个新元素
  • • 如果size是数组中的最后一个索引，说明堆已满，返回False；如果不是则继续
    • size先加1，因为size代表的是已经加入的元素个数，也就是现在堆的大小
    • 给heap数组赋值
    • 调用shift_up，按规则调整到二叉树中的合适位置
    • 加入成功，返回True
  • pop：弹出堆顶元素
  • • 堆顶元素在索引为1的位置，先用val保留该值
    • 将堆末尾元素heap[size]赋值给堆顶，然后将末尾位置置0
    • 将堆顶元素按规则下沉
    • 返回已暂存的堆顶元素val
  • peek：返回堆顶元素heap[1]
  • shift_up：将新加入到堆末尾的元素上浮到合适位置，因为是最小堆，这个合适位置也就是该位置以上节点均比它小
  • • 先用val暂存索引为$i$​位置的元素
    • 父节点的索引为 i >> 1，只要父节点索引不为0，就可以继续上浮
    • 如果当前元素val比父节点要小，和父节点交换位置，继续上浮
    • 一旦当前元素val比父节点大了，说明该位置以上的节点已经都比它小了，不用上浮了，break结束循环
    • 最后不要忘了把当前$i$​位置元素赋值为val
  • shift_down：将刚赋值为原来堆末尾元素的堆顶元素下沉
  • • 先用val暂存索引为$i$​的元素
    • 其左子节点索引为 i << 1
    • 如果左子节点并非是堆中最后一个元素，并且其右子节点的值比左子节点更小，索引 + 1，让父节点i和右子节点i << 1 | 1比较。
    • 父节点值val和子节点值比较，如果val更大，父节点和子节点交换位置，也就是父节点下沉
    • 一旦父节点值val比子节点小了，说明该位置以下的节点已经都比它大了，不用下沉了，break结束循环
    • 最后不要忘了把当前i位置元素赋值为val

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = Heap(k + 1)
        for num in nums:
            if not heap.push(num):
                heap.pop()
                heap.push(num)
        if heap.size == k + 1:
            heap.pop()
        return heap.peek()

class Heap:
    def __init__(self, length):
        self.heap = [0] * (length + 1)
        self.size = 0

    def push(self, val):
        if self.size == len(self.heap) - 1:
            return False
        self.size += 1
        self.heap[self.size] = val
        self.shift_up(self.size)
        return True

    def pop(self):
        val = self.heap[1]
        self.heap[1] = self.heap[self.size]
        self.heap[self.size] = 0
        self.size -= 1
        self.shift_down(1)
        return val

    def peek(self):
        return self.heap[1]

    def shift_up(self, i):
        val = self.heap[i]
        while i >> 1 > 0:
            parent = i >> 1
            if val < self.heap[parent]:
                self.heap[i] = self.heap[parent]
                i = parent
            else:
                break
        self.heap[i] = val

    def shift_down(self, i):
        val = self.heap[i]
        while i << 1 <= self.size:
            child = i << 1
            if child != self.size and self.heap[child + 1] < self.heap[child]:
                child += 1
            if val > self.heap[child]:
                self.heap[i] = self.heap[child]
                i = child
            else:
                break
        self.heap[i] = val

复杂度分析

• 时间复杂度：$O(nlogk)$ ，遍历数组$O(n)$​ ，上浮和下沉$O(logk)$
• 空间复杂度：$O(k)$

大顶堆

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        self.buildMaxHeap(nums)
        size = len(nums)
        for i in range(k-1):
            nums[0], nums[size-i-1] = nums[size-i-1], nums[0] #根节点跟末尾节点交换
            self.heapify(nums, 0, size-i-2)
        return nums[0]

    # 调整为大顶堆
    def heapify(self,nums, index, end):
        left = index * 2 + 1
        right = left + 1
        while left <= end:
            # 当前节点为非叶子结点
            max_index = index
            if nums[left] > nums[max_index]:
                max_index = left
            if right <= end and nums[right] > nums[max_index]:
                max_index = right
            if index == max_index:
                # 如果不用交换，则说明已经交换结束
                break
            nums[index], nums[max_index] = nums[max_index], nums[index]
            # 继续调整子树
            index = max_index
            left = index * 2 + 1
            right = left + 1
    # 初始化大顶堆
    def buildMaxHeap(self,nums):
        size = len(nums)
        # (size-2) // 2 是最后一个非叶节点，叶节点不用调整
        for i in range((size - 2) // 2, -1, -1):
            self.heapify(nums, i, size - 1)
        return nums

快速选择

方法讲解：

06快速排序算法

左右挖坑互填：

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        l = 0
        r = n - 1
        while True:
            idx = self.partition(nums, l, r)
            if idx == n-k:
                return nums[idx]
            elif idx < n-k:
                l = idx + 1
            else:
                r = idx - 1

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

左右交换：

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        l = 0
        r = n - 1
        while True:
            idx = self.partition(nums, l, r)
            if idx == n-k:
                return nums[idx]
            elif idx < n-k:
                l = idx + 1
            else:
                r = idx - 1

    def partition(self, nums, left, right):
        pivot = random.randint(left, right) #初始化一个待比较数据
        nums[left], nums[pivot] = nums[pivot], nums[left]
        pivot = nums[left]
        begin = left
        while left<right:
            while left<right and nums[right] >= pivot:#从后往前查找，直到找到一个比pivot更小的数
                right-=1
            while left<right and nums[left] <= pivot: #从前往后找，直到找到一个比pivot更大的数
                left +=1
            if left<right:
                nums[left], nums[right] = nums[right], nums[left]
        nums[begin], nums[left] = nums[left], nums[begin]
        return left
            
    

单方向遍历：

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        n = len(nums)
        l = 0
        r = n - 1
        while True:
            idx = self.partition(nums, l, r)
            if idx == n-k:
                return nums[idx]
            elif idx < n-k:
                l = idx + 1
            else:
                r = idx - 1

    def partition(self, nums, left, right):
        pivot = random.randint(left, right) #初始化一个待比较数据
        nums[left], nums[pivot] = nums[pivot], nums[left]
        pivot = nums[left]
        # pivot = nums[left]
        idx  = left
        for i in range(left + 1, right + 1):
            if nums[i] <= pivot:
                idx += 1
                nums[idx], nums[i] = nums[i], nums[idx]
        nums[idx], nums[left] = nums[left], nums[idx]
        return idx
    

参考

力扣（LeetCode） (leetcode-cn.com)