Skip to content

1962. Remove Stones to Minimize the Total 👍

  • Time: $O(k\log n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution {
 public:
  int minStoneSum(vector<int>& piles, int k) {
    int ans = accumulate(piles.begin(), piles.end(), 0);
    priority_queue<int> maxHeap;

    for (const int pile : piles)
      maxHeap.push(pile);

    for (int i = 0; i < k; ++i) {
      const int maxPile = maxHeap.top();
      maxHeap.pop();
      maxHeap.push(maxPile - maxPile / 2);
      ans -= maxPile / 2;
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
  public int minStoneSum(int[] piles, int k) {
    int ans = Arrays.stream(piles).sum();
    Queue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

    for (final int pile : piles)
      maxHeap.offer(pile);

    for (int i = 0; i < k; ++i) {
      final int maxPile = maxHeap.poll();
      maxHeap.offer(maxPile - maxPile / 2);
      ans -= maxPile / 2;
    }

    return ans;
  }
}

1
2
3
4
5
6
7
8
9
class Solution:
  def minStoneSum(self, piles: list[int], k: int) -> int:
    maxHeap = [-pile for pile in piles]
    heapq.heapify(maxHeap)

    for _ in range(k):
      heapq.heapreplace(maxHeap, maxHeap[0] // 2)

    return -sum(maxHeap)