99클럽 코테 스터디 22일차 TIL + 오늘의 학습 키워드 : 힙

오늘의 항해22일차 문제는
leetcode라는 사이트에서 2558번 문제

 
정답코드
:

import math
import heapq

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        # 최대 힙을 사용하기 위해 선물 수에 음수를 곱하여 저장
        max_heap = [-gift for gift in gifts]
        heapq.heapify(max_heap)

        for _ in range(k):
            # 최대 선물 수를 가진 더미 선택
            max_gift = -heapq.heappop(max_heap)
            
            # 남길 선물 수 (제곱근의 바닥)
            remaining = math.floor(math.sqrt(max_gift))
            
            # 남은 선물 수를 힙에 다시 추가
            heapq.heappush(max_heap, -remaining)

        # 남은 선물의 총 개수 계산
        remaining_total = -sum(max_heap)
        return remaining_total

코드 설명
1. 최대 힙 사용: heapq 모듈을 사용하여 최대 힙을 구현합니다. 파이썬의 heapq는 최소 힙만 지원하므로, 선물 수에 음수를 곱하여 최대 힙처럼 동작하게 합니다.
2. 반복문: k초 동안 반복하여 매번 최대 선물 수를 가진 더미를 선택합니다.
3. 제곱근 계산: 선택한 더미에서 남길 선물 수를 제곱근의 바닥으로 계산합니다.
4. 결과 계산: 모든 더미의 남은 선물 수를 합산하여 반환합니다.