Thomas Ngo

Software Engineer

3 minute read

Problem Statement

problem

Intuition

When we are allowed to perform ( k ) operations to reduce the size of the gift piles, we want to maximize the impact of each reduction. The operation transforms a pile of size ( x ) into a pile of size (\lfloor\sqrt{x}\rfloor). This means that applying the operation to a larger pile yields a larger overall reduction in that pile’s size compared to applying it to a smaller pile.

For example:

  • If ( x = 100 ), then (\lfloor\sqrt{100}\rfloor = 10), which is a big drop from 100 to 10.
  • If ( x = 9 ), then (\lfloor\sqrt{9}\rfloor = 3 ), which is a smaller drop from 9 to 3.

This suggests a “greedy” strategy: always pick the largest pile available to perform the operation on. By consistently reducing the largest pile, we achieve the greatest immediate decrease in total sum.

Approach

  1. Use a max-heap: To efficiently find and operate on the largest pile each time, we can use a max-heap (priority queue). Since Python’s heapq is a min-heap by default, we can invert the values (push negative values) to simulate a max-heap.

  2. Repeatedly pop the largest pile:
    • Extract the largest pile (which will be at the top of the heap).
    • Compute its new size by taking the floor of its square root.
    • Insert this new size back into the heap.
  3. Perform this operation ( k ) times:
    • After ( k ) iterations, no more operations can be performed.
  4. Compute the final sum:
    • After all operations are done, sum up all the piles in the heap to get the final answer.

This approach ensures that at each step, we achieve the maximum possible reduction by targeting the largest pile.

Complexity

  • Time Complexity: Every iteration requires:

    • Extracting the largest element from the heap: ( O(\log n) )
    • Inserting the modified element back into the heap: ( O(\log n) )

    Since we perform ( k ) such iterations, the time complexity is ( O(k \log n) ).

  • Space Complexity: We store all ( n ) piles in the heap, so the space complexity is ( O(n) ).

Code

import heapq
import math
from typing import List

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        # Create a max heap by inserting negative values
        max_heap = [-x for x in gifts]
        heapq.heapify(max_heap)

        # Perform k operations
        while k > 0:
            # Extract the largest pile
            largest = -heapq.heappop(max_heap)
            # Reduce the pile size
            new_size = math.floor(math.sqrt(largest))
            # Push the reduced pile back
            heapq.heappush(max_heap, -new_size)
            k -= 1

        # Compute the final sum
        return sum(-x for x in max_heap)

Editorial Solution

Approach 2: Sorted Array

class Solution:
    def pickGifts(self, gifts, k):
        n = len(gifts)

        # Create a copy of the gifts array and sort it
        sorted_gifts = sorted(gifts)

        # Perform the operation k times
        for _ in range(k):
            max_element = sorted_gifts[-1]
            sorted_gifts.pop()

            # Find the index where the square root of max_element should be inserted
            spot_of_sqrt = next(
                (
                    i
                    for i, value in enumerate(sorted_gifts)
                    if value > math.floor(math.sqrt(max_element))
                ),
                n,
            )

            # Insert the square root value at the correct position
            sorted_gifts.insert(
                spot_of_sqrt, math.floor(math.sqrt(max_element))
            )

        # Calculate the sum of the remaining elements in the sorted array
        number_of_remaining_gifts = sum(sorted_gifts)

        return number_of_remaining_gifts

Approach 3: Heap

class Solution:
    def pickGifts(self, gifts, k):
        # Create a max-heap from the 'gifts' array (negating values to simulate max-heap)
        gifts_heap = [-gift for gift in gifts]
        heapq.heapify(gifts_heap)

        # Perform the operation 'k' times
        for _ in range(k):
            # Get the maximum element from the heap (top element)
            max_element = -heapq.heappop(gifts_heap)

            # Insert the floor of the square root of the maximum element back into the heap
            heapq.heappush(gifts_heap, -math.floor(math.sqrt(max_element)))

        # Accumulate the sum of the elements in the heap
        number_of_remaining_gifts = 0
        while gifts_heap:
            number_of_remaining_gifts -= heapq.heappop(gifts_heap)

        return number_of_remaining_gifts

You may also enjoy