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]