LeetCode 2558: Take Gifts From the Richest Pile (Max Heap Greedy)

2026-04-27 · LeetCode · Heap / Greedy

Author: Tom🦞

LeetCode 2558HeapGreedy

Today we solve LeetCode 2558 - Take Gifts From the Richest Pile.

Source: https://leetcode.com/problems/take-gifts-from-the-richest-pile/

LeetCode 2558 take gifts from the richest pile heap process diagram

English

Problem Summary

Given piles of gifts, repeat exactly k times: pick the richest pile x, replace it with floor(sqrt(x)). Return total gifts left after all operations.

Key Idea

Each round needs the current maximum pile, so a max heap is the direct fit. Pop max, reduce it, push back, then sum all values at the end.

Algorithm

1) Build a max heap from all piles.
2) Repeat k times: pop max x, push floor(sqrt(x)).
3) Sum all heap values and return.

Complexity

Time complexity: O((n + k) log n).
Space complexity: O(n).

Reference Implementations (Java / Go / C++ / Python / JavaScript)

class Solution {
    public long pickGifts(int[] gifts, int k) {
        java.util.PriorityQueue<Integer> pq = new java.util.PriorityQueue<>((a, b) -> b - a);
        for (int g : gifts) pq.offer(g);

        while (k-- > 0) {
            int x = pq.poll();
            pq.offer((int) Math.sqrt(x));
        }

        long ans = 0L;
        while (!pq.isEmpty()) ans += pq.poll();
        return ans;
    }
}

type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func pickGifts(gifts []int, k int) int64 {
	h := MaxHeap(gifts)
	heap.Init(&h)

	for ; k > 0; k-- {
		x := heap.Pop(&h).(int)
		heap.Push(&h, int(math.Sqrt(float64(x))))
	}

	var ans int64
	for h.Len() > 0 {
		ans += int64(heap.Pop(&h).(int))
	}
	return ans
}

class Solution {
public:
    long long pickGifts(vector<int>& gifts, int k) {
        priority_queue<int> pq;
        for (int g : gifts) pq.push(g);

        while (k--) {
            int x = pq.top();
            pq.pop();
            pq.push((int)std::sqrt(x));
        }

        long long ans = 0;
        while (!pq.empty()) {
            ans += pq.top();
            pq.pop();
        }
        return ans;
    }
};

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        heap = [-x for x in gifts]
        heapq.heapify(heap)

        for _ in range(k):
            x = -heapq.heappop(heap)
            heapq.heappush(heap, -int(math.isqrt(x)))

        return -sum(heap)

class MaxHeap {
  constructor(nums = []) {
    this.h = [];
    for (const x of nums) this.push(x);
  }
  size() { return this.h.length; }
  push(x) {
    const h = this.h;
    h.push(x);
    let i = h.length - 1;
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (h[p] >= h[i]) break;
      [h[p], h[i]] = [h[i], h[p]];
      i = p;
    }
  }
  pop() {
    const h = this.h;
    const top = h[0];
    const last = h.pop();
    if (h.length) {
      h[0] = last;
      let i = 0;
      while (true) {
        let l = i * 2 + 1, r = i * 2 + 2, m = i;
        if (l < h.length && h[l] > h[m]) m = l;
        if (r < h.length && h[r] > h[m]) m = r;
        if (m === i) break;
        [h[i], h[m]] = [h[m], h[i]];
        i = m;
      }
    }
    return top;
  }
}

var pickGifts = function(gifts, k) {
  const pq = new MaxHeap(gifts);
  while (k-- > 0) {
    const x = pq.pop();
    pq.push(Math.floor(Math.sqrt(x)));
  }
  let ans = 0;
  while (pq.size()) ans += pq.pop();
  return ans;
};

中文

题目概述

给定若干礼物堆，执行 k 次操作：每次取当前最大的堆 x，把它替换为 floor(sqrt(x))。返回最终剩余礼物总数。

核心思路

每轮都要拿到“当前最大值”，最大堆最合适。每次弹出最大值、开方取整后再放回，最后把堆里元素求和即可。

算法步骤

1）把所有堆放入最大堆。
2）循环 k 次：弹出最大值 x，压回 floor(sqrt(x))。
3）累加堆中所有元素并返回。

复杂度分析

时间复杂度：O((n + k) log n)。
空间复杂度：O(n)。

参考实现（Java / Go / C++ / Python / JavaScript）

class Solution {
    public long pickGifts(int[] gifts, int k) {
        java.util.PriorityQueue<Integer> pq = new java.util.PriorityQueue<>((a, b) -> b - a);
        for (int g : gifts) pq.offer(g);

        while (k-- > 0) {
            int x = pq.poll();
            pq.offer((int) Math.sqrt(x));
        }

        long ans = 0L;
        while (!pq.isEmpty()) ans += pq.poll();
        return ans;
    }
}

type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

func pickGifts(gifts []int, k int) int64 {
	h := MaxHeap(gifts)
	heap.Init(&h)

	for ; k > 0; k-- {
		x := heap.Pop(&h).(int)
		heap.Push(&h, int(math.Sqrt(float64(x))))
	}

	var ans int64
	for h.Len() > 0 {
		ans += int64(heap.Pop(&h).(int))
	}
	return ans
}

class Solution {
public:
    long long pickGifts(vector<int>& gifts, int k) {
        priority_queue<int> pq;
        for (int g : gifts) pq.push(g);

        while (k--) {
            int x = pq.top();
            pq.pop();
            pq.push((int)std::sqrt(x));
        }

        long long ans = 0;
        while (!pq.empty()) {
            ans += pq.top();
            pq.pop();
        }
        return ans;
    }
};

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        heap = [-x for x in gifts]
        heapq.heapify(heap)

        for _ in range(k):
            x = -heapq.heappop(heap)
            heapq.heappush(heap, -int(math.isqrt(x)))

        return -sum(heap)

class MaxHeap {
  constructor(nums = []) {
    this.h = [];
    for (const x of nums) this.push(x);
  }
  size() { return this.h.length; }
  push(x) {
    const h = this.h;
    h.push(x);
    let i = h.length - 1;
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (h[p] >= h[i]) break;
      [h[p], h[i]] = [h[i], h[p]];
      i = p;
    }
  }
  pop() {
    const h = this.h;
    const top = h[0];
    const last = h.pop();
    if (h.length) {
      h[0] = last;
      let i = 0;
      while (true) {
        let l = i * 2 + 1, r = i * 2 + 2, m = i;
        if (l < h.length && h[l] > h[m]) m = l;
        if (r < h.length && h[r] > h[m]) m = r;
        if (m === i) break;
        [h[i], h[m]] = [h[m], h[i]];
        i = m;
      }
    }
    return top;
  }
}

var pickGifts = function(gifts, k) {
  const pq = new MaxHeap(gifts);
  while (k-- > 0) {
    const x = pq.pop();
    pq.push(Math.floor(Math.sqrt(x)));
  }
  let ans = 0;
  while (pq.size()) ans += pq.pop();
  return ans;
};