• Difficulty: Medium
• Patterns:
  • Heap

---

Problem Definition

You are given an array of magic bags containing different numbers of candies. It takes you 1 minute to eat a full bag of candies regardless of the number of candies in the bag. Once you finish eating the candies, the bag is magically refilled so that if the bag initially contained x candies then at the end of the minute, the bag is refilled to contain floor(x / 2) candies. Determine how many candies you can eat in $k$ minutes.

---

Examples

Some test cases to verify our solution.

assert candies([2, 1, 7, 4, 2], 3) == 14 # because => 7 [2, 1, 3, 4, 2] => 7+4 [2, 1, 3, 2, 2] =>
    # 7+4+3 [2, 1, 1, 2, 2]
assert candies([19, 78, 76, 72, 48, 8, 24, 74, 29], 3) == 228

---

Brute-Force Solution

We will sort the array of candy bags and then eat from the bag with the most candies. After that, we will put the bag back into the array of bags and sort the whole array.

def candies(arr, k):
    l = sorted(arr)
    candies = 0
    for _ in range(k):
        bag = l.pop()
        candies += bag
        l.append(bag // 2)
        l = sorted(l)
    return candies

---

Brute-Force Solution Complexity

Let $n = $ len(arr).

Runtime complexity should be $O(k \cdot n \cdot log(n))$. Our for loop has $k$ iterations, and in each iteration, we are sorting an array of length $n$, which results in a nested complexity of $O(n \cdot log(n))$. Auxiliary space complexity is $O(n)$ because we are creating a copy of arr.

---

Improved Solution

We can remove the $log(n)$ factor by traversing the array at every iteration to find the largest element.

def candies(arr, k):
    candies = 0
    for _ in range(k):
        bag = max(arr)
        arr.remove(bag)
        candies += bag
        arr.append(bag // 2)
    return candies

---

Improved Solution Complexity

Let $n = $ len(arr).

Runtime complexity should be $O(k \cdot n)$. Our for loop has $k$ iterations, and in each iteration, we are traversing an array of length $n$, which results in a nested complexity of $O(n)$.

Auxiliary space complexity is $O(1)$ but we are destroying arr.

---

Heap Solution

We can use a max-heap to simplify our code and improve runtime complexity.

import heapq

def candies(arr, k):
    heapq._heapify_max(arr) # nlogn
    candies = 0
    for _ in range(k): # k
        bag = arr[0]
        candies += bag
        heapq._heapreplace_max(arr, bag // 2) # logn
    return candies

---

Heap Solution Complexity

Let $n = $ len(arr).

Runtime complexity should be $O(n \cdot log(n) + k \cdot log(n)) = O((n + k) \cdot log(n))$. The construction of the heap has a complexity of $O(n \cdot log(n))$. Our for loop has $k$ iterations, and in each iteration, we are doing a pop_push into a heap of size $n$, resulting in a nested complexity of $O(log(n))$.

Auxiliary space complexity is $O(1)$ because we are modifying arr in place.