[Coding test] 9. Heap

whitehousechef·2025년 3월 3일

Heap

So heap doesnt necessarily store elements in a sorted order, contrary to my belief. The only sorting thing is at its root (i.e. value at 0th index)

min heap

Example

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

Finding kth largest number in a given list means we have to make a min heap of size k and the 0th index will store the answer.

initial

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

but theres so many if and else and whatever. What if we just make a heap of size k initially and pop whenever a given number is larger than the root stored in our heap?

sol

    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = nums[:k]  # Take first k elements
        heapq.heapify(heap)  # Convert them into a min-heap (O(k))

        for num in nums[k:]:  # Process remaining elements
            heapq.heappushpop(heap, num)  # Push num and pop smallest in one step

        return heap[0]  # Root of the heap is the k-th largest element

import heapq
from typing import List

def findKthLargest(nums: List[int], k: int) -> int:
    # Create a min-heap with the first k elements
    heap = nums[:k]
    heapq.heapify(heap)  # Convert the first k elements into a min-heap

    # Process the remaining elements
    for num in nums[k:]:
        if num > heap[0]:  # If the current number is larger than the smallest in the heap
            heapq.heappop(heap)  # Remove the smallest element
            heapq.heappush(heap, num)  # Add the current element

    # The root of the heap will be the k-th largest element
    return heap[0]