Tech for good

[Leetcode/Heap] 2558. Take Gifts From the Richest Pile 본문

IT/Computer Science

[Leetcode/Heap] 2558. Take Gifts From the Richest Pile

Diana Kang 2025. 4. 4. 22:57

Example 1:
gifts = [25, 64, 9, 4, 100], k = 4

[25, 64, 9, 4, 100]
-> (1회) [25, 64, 9, 4, 10]
-> (2회) [25, 8, 9, 4, 10]
-> (3회) [5, 8, 9, 4, 10]
-> (4회) [5, 8, 9, 4, 3] (floor or the square root - floor: 내림)

5 + 8 + 9 + 4 + 3 = 29

 


Follow these steps:

  1. Pick the largest pile of gifts.
  2. Replace it with the floor of its square root.
  3. Repeat this for k seconds.

🔧 Best Tool for This?

We need to efficiently get the largest number repeatedly — that's a classic use case for a max heap.

BUT: Python’s heapq module only supports min-heaps, so we simulate a max-heap by inserting negative numbers.

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        # Step 1: Convert to max-heap by inserting negative values
        max_heap = [-g for g in gifts]
        heapq.heapify(max_heap)

        # Step 2: Perform the operation k times
        for _ in range(k):
            # Pop the largest (most negative = smallest negative)
            largest = -heapq.heappop(max_heap)
            reduced = math.isqrt(largest)  # floor(sqrt(x))
            heapq.heappush(max_heap, -reduced)

        # Step 3: Sum up the remaining gifts (convert back to positive)
        return -sum(max_heap)

💡 Notes:

  • math.isqrt(x) is a fast and accurate way to do floor(sqrt(x)) for integers.
  • Using a heap makes this efficient even when the list is long or k is big.

📚 math.isqrt(x) explained:

math.isqrt(x) is a function in Python that returns the integer (floor) square root of x.

It gives you the largest integer whose square is less than or equal to x.

import math

print(math.isqrt(10))   # 3, because 3^2 = 9 (and 4^2 = 16 > 10)
print(math.isqrt(25))   # 5, because 5^2 = 25
print(math.isqrt(100))  # 10
print(math.isqrt(1))    # 1