NeetCode 150 [Heap/Priority Queue]medium:KthLargestElementInAnArray

NeetCodeのSolutionを書いていく

Kth Largest Element in an Array

問題

ソートされていない整数の配列numsと整数kが与えられます。
k番目に大きな要素を返しなさい。
ソートせずに解けますか？

例

Input: nums = [2,3,1,5,4], k = 2
Output: 4

Input: nums = [2,3,1,1,5,5,4], k = 3
Output: 4

制約

• 1 <= k <= nums.length <= 10000
• -1000 <= nums[i] <= 1000
• 計算量：O(n*longk)
• メモリ：O(k)

メモ

pythonのソートはO(n log n)の計算量がかかるらしい。(by ChatGPT)
これだと指定のO(n log k)を超えてしまう。
しかもソートしないで解けますか？とのこと。
メモリがO(k)何だから、k個までは保持できる。
heapqのにpushしながら、小さい方は消していったらどうだろうか。最後にright popして返す。
これだとheapqに保持する数は最大kなのでpushに最大でO(log k)となり、n個処理するのにO(n log k)の時間がかかってしまうでいけそう。

import heapq
from typing import List


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

以下で行けたらしい。

class Solution:
    def findKthLargest(self, nums, k):
        return heapq.nlargest(k, nums)[-1]