Week 3-2. [Heap] 215. Kth Largest Element in an Array (Python3)

문제 해결 과정

✍️ 문제 요약

https://leetcode.com/problems/kth-largest-element-in-an-array/?envType=study-plan-v2&envId=top-interview-150 

Kth Largest Element in an Array - LeetCode

➡️ 주어진 배열을 정렬하지 않고 k번째로 큰 수를 반환하라.

 

🔎 어떤 식으로 접근했는가?

정렬을 하지 말라는 요구사항이 있었으므로 최대힙을 사용하자.

 

우선 파이썬에서 제공하는 힙 라이브러리를 사용해 풀어보고, 힙을 직접 구현해서도 풀어볼 것이다.

파이썬에서 제공하는 힙 라이브러리는 최소힙이기 때문에 마이너스 기호를 이용해 최대힙처럼 동작하게 할 수 있다.

 

힙을 직접 구현할 때는 push와 pop을 구현해야 하는데, 각각에서 사용되는 bubbleUp과 bubbleDown 로직이 핵심이다.

 

🎯 풀이

파이썬 라이브러리를 사용한 풀이이다.

import heapq as hq
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        lst = [] # 힙을 모방할 리스트

        for x in nums:
            hq.heappush(lst, -x) # 최대힙

        ans = 0
        for x in range(k):
            ans = -hq.heappop(lst)
        return ans

 

힙을 직접 구현하여 사용한 풀이이다.

class Heap:
    def __init__(self):
        self.array = []
        self.numberOfItems = 0
    
    def push(self, e: int) -> None:
        self.array.append(e) # 새로운 요소를 맨 마지막에 저장
        self.bubbleUp(self.numberOfItems) # Heap의 맨 마지막 요소인 array[numberOfItems]부터 bubble Up!
        self.numberOfItems += 1

    def bubbleUp(self, i: int) -> None:
        child, parent = i, (i-1)//2
        if parent >= 0 and self.array[child] > self.array[parent]:
            self.array[child], self.array[parent] = self.array[parent], self.array[child] # swap
            self.bubbleUp(parent)
    
    def pop(self) -> int:
        mx = self.array[0]
        self.array[0] = self.array[self.numberOfItems-1] # 마지막 요소를 0번째 인덱스에 저장
        self.numberOfItems -= 1
        self.bubbleDown(0) # # Heap의 첫 번째 요소인 array[0]부터 bubble Down!
        return mx

    def bubbleDown(self, i: int) -> None: # i는 부모 인덱스
        left, right = 2*i+1, 2*i+2 # 왼쪽 자식, 오른쪽 자식
        if left <= self.numberOfItems - 1:
            # 둘 중 더 큰 자식 선택
            if right <= self.numberOfItems and self.array[right] > self.array[left]:
                left = right
            # 위에서 선택된 자식이 부모보다 크다면 swap
            if self.array[i] < self.array[left]:
                self.array[i], self.array[left] = self.array[left], self.array[i] # swap
                self.bubbleDown(left)

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = Heap()
        for x in nums:
            heap.push(x)
        ans = 0
        for i in range(k):
            ans = heap.pop()
        return ans

 

다른 사람의 코드를 탐구해보기

파이썬 라이브러리를 이용한 풀이이며 heapify()를 사용해 기존 nums 리스트를 최소힙으로 변환하여 이용한다. 

import heapq as hq
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        hq.heapify(nums) # nums를 최소힙으로 변환
        amount = len(nums)
        
        while amount > k:
            hq.heappop(nums)
            amount -= 1
            
        return nums[0]