[Priority Queue]99클럽 코테 스터디 22일차 TIL + 2692-take-gifts-from-the-richest-pile

2692. Take Gifts From the Richest Pile

Easy

---

You are given an integer array gifts denoting the number of gifts in various piles. Every second, you do the following:

• Choose the pile with the maximum number of gifts.
• If there is more than one pile with the maximum number of gifts, choose any.
• Leave behind the floor of the square root of the number of gifts in the pile. Take the rest of the gifts.

Return the number of gifts remaining after k seconds.

 

Example 1:

Input: gifts = [25,64,9,4,100], k = 4
Output: 29
Explanation: 
The gifts are taken in the following way:
- In the first second, the last pile is chosen and 10 gifts are left behind.
- Then the second pile is chosen and 8 gifts are left behind.
- After that the first pile is chosen and 5 gifts are left behind.
- Finally, the last pile is chosen again and 3 gifts are left behind.
The final remaining gifts are [5,8,9,4,3], so the total number of gifts remaining is 29.

Example 2:

Input: gifts = [1,1,1,1], k = 4
Output: 4
Explanation: 
In this case, regardless which pile you choose, you have to leave behind 1 gift in each pile. 
That is, you can't take any pile with you. 
So, the total gifts remaining are 4.

 

Constraints:

• 1 <= gifts.length <= 103
• 1 <= gifts[i] <= 109
• 1 <= k <= 103

---

<내 코드>

import heapq
import math

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        answer = 0

        # 1. gits 를 maxheap 으로 변환
        max_heap = [-x for x in gifts]
        heapq.heapify(max_heap)

        # 2. K 번 반복하기
        for _ in range(k):
            # (1) heapop 후 제곱근 (내림)
            if not max_heap:
                break
            piles = - heapq.heappop(max_heap)
            print(piles)

            # 1) 만약 1이라면 제곱근 안하기
            if piles == 1:
                answer += 1
                continue
            # 2) 제곱근
            squ_pile = int(math.sqrt(piles))
            # (2) heappush
            heapq.heappush(max_heap, -squ_pile)

        # 3. list 의 값 더하기
        answer += (-sum(max_heap))

        return answer

시간복잡도 : O(k log n)

 

- answer 을 딱히 누계하면서 생각 할 필요는 없음

- 생각해보면 for 를 요새 자주쓰는데 while 을 생각해보자

 

<모범 사례>

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        nums = [-num for num in gifts]
        heapify(nums)
        
        while k:
            tmp = math.isqrt(-heappop(nums))
            heappush(nums, -tmp)
            k -= 1
            
        return -sum(nums)

- 다음과 같이 그냥 쉽게 해주면 됨

 

 

<보충학습>

heap : log 시간 내 최대,최소 요소의 효율적 검색이 가능함 / 대용량 데이터 반복에 중요