Problem of The Day: Maximize Happiness of Selected Children

Problem Statement

Heap Approach - TLE

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        max_heap = []
        for happy_val in happiness:
            heapq.heappush(max_heap, happy_val * -1)

        res = []
        while max_heap and len(res) < k:
            res.append(heapq.heappop(max_heap) * -1)
            curr_heap = []
            while max_heap:
                val = (heapq.heappop(max_heap) * -1) - 1
                if val > 0:
                    heapq.heappush(curr_heap, val * -1)

            max_heap = curr_heap

        return sum(res)

Improved Heap Approach - passed

Intuition

My thought here is to maximize the total happiness sum by selecting the highest happiness values from the input list. Since we can only choose k values, we need to prioritize the highest ones.

Approach

To implement this, I’ll utilize a max heap to efficiently extract the highest happiness values. I’ll push the negated values into the heap to simulate a max heap. Then, I’ll iterate through the heap, popping elements and subtracting an index value from them to ensure we prioritize the earlier elements. We’ll continue this process until we have selected k values or the heap is empty.

Complexity

• Time complexity: Pushing all elements onto the heap takes O(n _ log(n)), where n is the number of elements in the input list. Popping elements and constructing the result list takes O(k _ log(n)), as we perform k pop operations. Overall, the time complexity is O(n _ log(n) + k _ log(n)).

• Space complexity: O(n) for storing the heap and the result list.

Code

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        max_heap = []
        for happy_val in happiness:
            heapq.heappush(max_heap, happy_val * -1)

        res = []
        i = 0
        while max_heap and len(res) < k:
            val = (heapq.heappop(max_heap) * -1) - i
            if val > 0:
                res.append(val)
            i += 1

        return sum(res)

Editorial Solution

Approach 1: Sort + Greedy

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        # Sort in descending order
        happiness.sort(reverse=True)

        total_happiness_sum = 0
        turns = 0

        # Calculate the maximum happiness sum
        for i in range(k):
            # Adjust happiness and ensure it's not negative
            total_happiness_sum += max(happiness[i] - turns, 0)

            # Increment turns for the next iteration
            turns += 1

        return total_happiness_sum

• Time: O(nlogn)
• Space: O(n)

Approach 2: Max Heap / Priority Queue + Greedy

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        # Convert the list into a max heap by inverting the happiness values
        # Python's default heap data structure is a min heap
        max_heap = [-h for h in happiness]
        heapq.heapify(max_heap)

        total_happiness_sum = 0
        turns = 0

        for i in range(k):
            # Invert again to get the original value
            total_happiness_sum += max(-heapq.heappop(max_heap) - turns, 0)

            # Increment turns for the next iteration
            turns += 1

        return total_happiness_sum

• Time: O(nlogn + klogn)
• Space: O(n)