Skip to content

2558. Take Gifts From the Richest Pile 👍

  • 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
20
21
class Solution {
 public:
  long long pickGifts(vector<int>& gifts, int k) {
    long long ans = 0;
    priority_queue<int> maxHeap;

    for (const int gift : gifts)
      maxHeap.push(gift);

    for (int i = 0; i < k; ++i) {
      const int squaredMax = sqrt(maxHeap.top());
      maxHeap.pop();
      maxHeap.push(squaredMax);
    }

    while (!maxHeap.empty())
      ans += maxHeap.top(), maxHeap.pop();

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution {
  public long pickGifts(int[] gifts, int k) {
    long ans = 0;
    Queue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

    for (final int gift : gifts)
      maxHeap.offer(gift);

    for (int i = 0; i < k; ++i) {
      final int squaredMax = (int) Math.sqrt(maxHeap.poll());
      maxHeap.offer(squaredMax);
    }

    while (!maxHeap.isEmpty())
      ans += maxHeap.poll();

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
  def pickGifts(self, gifts: List[int], k: int) -> int:
    maxHeap = [-gift for gift in gifts]
    heapq.heapify(maxHeap)

    for _ in range(k):
      squaredMax = int(sqrt(-heapq.heappop(maxHeap)))
      heapq.heappush(maxHeap, -squaredMax)

    return -sum(maxHeap)