[LeetCode] 2558. Take Gifts From the Richest Pile (Python)

 

난이도: 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 cant take any pile with you.
So, the total gifts remaining are 4.

 

제한 사항

---

• 1 <= gifts.length <= 10^3
• 1 <= gifts[i] <= 10^9
• 1 <= k <= 10^3

 

✏️Solution(솔루션)

---

class Solution(object):
    def pickGifts(self, gifts, k):
        gifts = [-gift for gift in gifts]
        heapq.heapify(gifts)

        for _ in range(k):
            max_gift = heapq.heappop(gifts)
            heapq.heappush(gifts, -floor(sqrt(-max_gift)))
        
        return int(-sum(gifts))

 

문제만 잘 이해하면 별로 어렵지 않은 문제인 것 같다.

 

결국에는 gifts에서 max값을 뽑아 square root를 취해 내림하는데 이 작업을 k번 하면된다.

 

모든 작업이 끝났을 때 gifts의 전체 요소 합을 return 하면 정답을 맞출 수 있다.

 

풀 수 있는 다양한 방법이 있겠지만, 나는 heapq를 사용해 제일 큰 값을 뽑고 그 값에 필요한 연산을 취한다음 다시 heapq에 넣는 방식으로 문제를 해결했다.

 

---

2558. Take Gifts From the Richest Pile