LeetCode 2233 - Maximum Product After K Increments

Difficulty: 🟡 Medium
Topics: Array, Greedy, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to maximize the product of an array of non-negative integers after performing at most k increment operations, where each operation increases any single element by 1. In other words, given an array nums and an integer k, we want to distribute these k increments across the elements of nums in such a way that the resulting product of all elements is as large as possible. Finally, we return the product modulo 10^9 + 7 because the numbers involved can become very large.

The input nums is an array of non-negative integers, and k is a positive integer up to 10^5. The array length can also be up to 10^5, and the individual elements can be as large as 10^6. This indicates that a brute-force approach that simulates each possible increment would be computationally infeasible because k could be very large and the array could have up to 100,000 elements.

Edge cases to consider include arrays where all elements are zero, arrays with very large elements, or cases where k is larger than the difference needed to equalize all elements. The problem guarantees non-negative integers, so we do not need to handle negative values.

Approaches

The brute-force approach would involve incrementing one element at a time, always trying every possible choice of which element to increment. At each step, we would compute the product and keep track of the maximum. This approach is correct because eventually all possible distributions of increments are considered, but it is too slow for the input limits, as it would take O(k * n) time.

The optimal approach relies on a key insight: to maximize the product of numbers, we should always increment the smallest element in the array. This is because the product grows faster when smaller numbers increase than when larger numbers increase, given the properties of multiplication. Using a min-heap (priority queue) allows us to efficiently extract the smallest element, increment it, and reinsert it, repeating this process k times. This ensures the operations are performed in O(log n) time per increment, resulting in O(k log n) time overall, which is feasible for the constraints.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(k * n)	O(1)	Increment smallest element by checking all elements each time
Optimal	O(k log n)	O(n)	Use a min-heap to efficiently find and increment the smallest element

Algorithm Walkthrough

Insert all elements of nums into a min-heap. This allows us to quickly identify and extract the smallest element at each step.
Repeat the following k times:
Extract the smallest element from the heap.
Increment this element by 1.
Push the incremented value back into the heap.
After performing all k increments, compute the product of all elements remaining in the heap.
Return the product modulo 10^9 + 7 to prevent overflow.

Why it works: By always incrementing the smallest element, we are maximizing the growth of the product. This strategy is guaranteed to produce the maximum product because increasing smaller numbers has a higher multiplicative impact than increasing larger numbers, which follows from the arithmetic-geometric mean inequality.

Python Solution

from typing import List
import heapq

class Solution:
    def maximumProduct(self, nums: List[int], k: int) -> int:
        MOD = 10**9 + 7
        heapq.heapify(nums)
        
        for _ in range(k):
            smallest = heapq.heappop(nums)
            heapq.heappush(nums, smallest + 1)
        
        result = 1
        for num in nums:
            result = (result * num) % MOD
        
        return result

The code first converts nums into a min-heap, which allows efficient retrieval of the smallest element. The for loop applies k increments, always to the smallest element. Finally, the product of the array is computed modulo 10^9 + 7 to handle large numbers.

Go Solution

package main

import (
    "container/heap"
)

func maximumProduct(nums []int, k int) int {
    const MOD = 1_000_000_007
    
    h := IntHeap(nums)
    heap.Init(&h)
    
    for i := 0; i < k; i++ {
        smallest := heap.Pop(&h).(int)
        heap.Push(&h, smallest+1)
    }
    
    result := 1
    for _, num := range h {
        result = (result * num) % MOD
    }
    
    return result
}

type IntHeap []int

func (h IntHeap) Len() int            { return len(h) }
func (h IntHeap) Less(i, j int) bool  { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *IntHeap) Push(x any)         { *h = append(*h, x.(int)) }
func (h *IntHeap) Pop() any {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

Go-specific differences: The heap operations are slightly more verbose, requiring a custom IntHeap type that implements heap.Interface. Overflow is handled by modular arithmetic, similar to Python. Nil slices are not a concern here because the input guarantees at least one element.

Worked Examples

Example 1: nums = [0,4], k = 5

Step	Heap State	Operation
0	[0,4]	Initial heap
1	[1,4]	Increment 0
2	[2,4]	Increment 1
3	[3,4]	Increment 2
4	[4,4]	Increment 3
5	[5,4]	Increment 4

Final product: 5 * 4 = 20.

Example 2: nums = [6,3,3,2], k = 2

Step	Heap State	Operation
0	[2,3,3,6]	Initial heap
1	[3,3,3,6]	Increment 2
2	[3,4,3,6]	Increment 3

Final product: 3 * 4 * 3 * 6 = 216.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(k log n)	Each increment requires `O(log n)` for heap pop and push, repeated `k` times
Space	O(n)	The min-heap stores all `n` elements

The algorithm is efficient for large k and n because heap operations are logarithmic and the product computation is linear.

Test Cases

# Provided examples
assert Solution().maximumProduct([0,4], 5) == 20  # Increment smallest element 5 times
assert Solution().maximumProduct([6,3,3,2], 2) == 216  # Increment smallest two elements

# Edge cases
assert Solution().maximumProduct([0], 1) == 1  # Single element
assert Solution().maximumProduct([1,1,1], 3) == 8  # Equal elements, multiple increments
assert Solution().maximumProduct([1000000], 100000) == 1100000  # Large element with large k
assert Solution().maximumProduct([0,0,0], 3) == 1  # All zeros
assert Solution().maximumProduct([2,2,2], 0) == 8  # No increments

Test	Why
[0,4], 5	Standard small example from problem
[6,3,3,2], 2	Standard example with multiple increments
[0], 1	Single element increment
[1,1,1], 3	Equal elements with multiple increments
[1000000], 100000	Large element stress test
[0,0,0], 3	All zeros edge case
[2,2,2], 0	No increments allowed

Edge Cases

Single-element array: When nums contains only one element, all k increments apply to that element. The algorithm correctly handles this because the heap will contain just one element, and repeated increments are applied directly.
All zeros: If all elements are zero, repeated increments must raise one element at a time. The heap ensures the smallest element (always zero initially) is incremented first, producing the correct maximum product without division by zero or negative numbers.
Large k relative to array size: If k is much larger than the number of elements, increments may repeatedly be applied to the same element. The heap ensures that increments always target the smallest element, efficiently redistributing increments even in this extreme case.
Large numbers: If nums contains very large numbers, computing the product can overflow