[LeetCode] 1962. Remove Stones to Minimize the Total (Python)

 

난이도: Medium

 

문제 설명

---

You are given a 0-indexed integer array piles, where piles[i] represents the number of stones in the ith pile, and an integer k. You should apply the following operation exactly k times:

 

• Choose any piles[i] and remove floor(piles[i] / 2) stones from it.

Notice that you can apply the operation on the same pile more than once.

Return the minimum possible total number of stones remaining after applying the k operations.

floor(x) is the greatest integer that is smaller than or equal to x (i.e., rounds x down).

 

문제 예제

---

Example 1:

Input: piles = [5,4,9], k = 2
Output: 12
Explanation: Steps of a possible scenario are:
- Apply the operation on pile 2. The resulting piles are [5,4,5].
- Apply the operation on pile 0. The resulting piles are [3,4,5].
The total number of stones in [3,4,5] is 12.

 

Example 2:

Input: piles = [4,3,6,7], k = 3
Output: 12
Explanation: Steps of a possible scenario are:
- Apply the operation on pile 2. The resulting piles are [4,3,3,7].
- Apply the operation on pile 3. The resulting piles are [4,3,3,4].
- Apply the operation on pile 0. The resulting piles are [2,3,3,4].
The total number of stones in [2,3,3,4] is 12.

 

제한사항

---

• 1 <= piles.length <= 10^5
• 1 <= piles[i] <= 10^4
• 1 <= k <= 10^5

 

✏️ Solution(솔루션)

class Solution:
    def minStoneSum(self, piles: List[int], k: int) -> int:
        piles_with_idx = [(-pile, idx) for idx, pile in enumerate(piles)]
        heapq.heapify(piles_with_idx)

        for _ in range(k):
            _, idx = heapq.heappop(piles_with_idx)
            piles[idx] -= piles[idx]//2
            heapq.heappush(piles_with_idx, (-piles[idx], idx))
        
        return sum(piles)

 

heapq를 이용해 문제를 해결했다.

 

먼저 pile의 값과 idx를 tuple로 만들어 piles_with_idx라는 리스트를 하나 만들고 이를 heapq로 사용하기 위해 heapify를 진행했다. 이때 큰 값을 먼저 pop하기 위해 pile에 -시켰다.

 

다음으로, piles_with_idx에서 heappop을 진행하고 기존 piles[idx]에 값을 변경해주었다. 연산된 값은 다시 piles_with_idx에 넣어주었고, 이러한 작업을 k번 동안 반복하였다.

 

모든 연산이 끝나면 piles의 합을 구해 return하여 정답을 맞출 수 있었다.

 

---

1962. Remove Stones to Minimize the Total