Teaching Kids Programming – Algorithms to Take Max/Min, Square Root, and Compute Sum after K times (Heap, Sort, Index)

Teaching Kids Programming: Videos on Data Structures and Algorithms

Take Gifts From the Richest Pile

You are given an integer array gifts denoting the number of gifts in various piles. Every second, you do the following:

Choose the pile with the maximum number of gifts.
If there is more than one pile with the maximum number of gifts, choose any.
Leave behind the floor of the square root of the number of gifts in the pile. Take the rest of the gifts.
Return the number of gifts remaining after k seconds.

Example 1:
Input: gifts = [25,64,9,4,100], k = 4
Output: 29
Explanation:
The gifts are taken in the following way:
– In the first second, the last pile is chosen and 10 gifts are left behind.
– Then the second pile is chosen and 8 gifts are left behind.
– After that the first pile is chosen and 5 gifts are left behind.
– Finally, the last pile is chosen again and 3 gifts are left behind.
The final remaining gifts are [5,8,9,4,3], so the total number of gifts remaining is 29.
Example 2:

Input: gifts = [1,1,1,1], k = 4
Output: 4
Explanation:
In this case, regardless which pile you choose, you have to leave behind 1 gift in each pile.
That is, you can’t take any pile with you.
So, the total gifts remaining are 4.

Constraints:
1 <= gifts.length <= 10^3
1 <= gifts[i] <= 10^9
1 <= k <= 10^3

Hints:
How can you keep track of the largest gifts in the array
What is an efficient way to find the square root of a number?
Can you keep adding up the values of the gifts while ensuring they are in a certain order?
Can we use a priority queue or heap here?

Algorithms to Take Max/Min from Numbers Array and Push Square Root Back

This can be done in 3 approaches: sorting, find the index, and using a heap.

Sorting the Numbers

We can do what it says, by simulating the process. We can sort the numbers to get the maximum. The time complexity is O(K*N*LogN+N) if N is the number of elements in the array.

1
2
3
4
5
6

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        for _ in range(k):
            A.sort()
            A[-1] = isqrt(A[-1])
        return sum(A)

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        for _ in range(k):
            A.sort()
            A[-1] = isqrt(A[-1])
        return sum(A)

The time complexity for accumulating the sums of the numbers is O(N) but it is trivial to the O(K*N*LogN) and thus we dropped this term. The space complexity for sorting is considered to be O(1) although in practice, by using Quick Sort, it could be O(LogN).

We use math.isqrt function to get the integer square root, which is the same as int(math.sqrt).

Find the Max Index

We don’t need to sort the array just to get the maximum or minmum value as sorting requires O(NLogN). Instead we can do the linear scan to find the max/min.

1
2
3
4
5
6
7

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        for _ in range(k):
            mx = max(A)
            ix = A.index(mx)
            A[ix] = isqrt(A[ix])
        return sum(A)

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        for _ in range(k):
            mx = max(A)
            ix = A.index(mx)
            A[ix] = isqrt(A[ix])
        return sum(A)

We use max function to get the max value, then we use find or index to locate the index of the max value, and then, we square root the maximum value. The time complexity is O(K*N+N) which is O(K*N)

Building a Heap and Heap Pop

We can build a heap data structure from N elements, which takes O(N). Note that if we build the heap data structure (max heap or min heap) by heappush one element at a time, this will take O(NLogN) instead.

The time complexity for using a heap is O(N+K*LogN+N) which is O(N+K*LogN), which yields the optimal time complexity among the three approaches described in this post. The space complexity can be O(1) constant if we transforming the original array to negatives rather than allocating for a separate array.

1
2
3
4
5
6
7
8

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        A = [-i for i in A]
        heapify(A)
        for _ in range(k):
            x = -heappop(A)
            heappush(A, -isqrt(x))
        return -sum(A)

class Solution:
    def pickGifts(self, A: List[int], k: int) -> int:
        A = [-i for i in A]
        heapify(A)
        for _ in range(k):
            x = -heappop(A)
            heappush(A, -isqrt(x))
        return -sum(A)

–EOF (The Ultimate Computing & Technology Blog) —