Maximum Product After K Increments

Med
#2034Time: O((N+k) log N). Building the heap from N elements takes O(N log N) time. Then, we perform `k` operations, each involving a `poll` and an `add`, both of which take O(log N) time. This part is O(k log N). The final product calculation involves polling N elements, taking O(N log N). The total complexity is dominated by these operations.Space: O(N), as the `PriorityQueue` needs to store all N elements from the input array.1 company
Patterns
Companies

Prompt

You are given an array of non-negative integers nums and an integer k. In one operation, you may choose any element from nums and increment it by 1.

Return the maximum product of nums after at most k operations. Since the answer may be very large, return it modulo 109 + 7. Note that you should maximize the product before taking the modulo. 

 

Example 1:

Input: nums = [0,4], k = 5
Output: 20
Explanation: Increment the first number 5 times.
Now nums = [5, 4], with a product of 5 * 4 = 20.
It can be shown that 20 is maximum product possible, so we return 20.
Note that there may be other ways to increment nums to have the maximum product.

Example 2:

Input: nums = [6,3,3,2], k = 2
Output: 216
Explanation: Increment the second number 1 time and increment the fourth number 1 time.
Now nums = [6, 4, 3, 3], with a product of 6 * 4 * 3 * 3 = 216.
It can be shown that 216 is maximum product possible, so we return 216.
Note that there may be other ways to increment nums to have the maximum product.

 

Constraints:

  • 1 <= nums.length, k <= 105
  • 0 <= nums[i] <= 106

Approaches

3 approaches with complexity analysis and trade-offs.

The brute-force approach is slow because it repeatedly scans the entire array to find the minimum element. We can optimize this process by using a data structure that provides fast access to the minimum element. A min-heap (implemented as a PriorityQueue in Java) is perfect for this task.

Algorithm

  • Create a PriorityQueue (min-heap) and add all elements from nums into it.
  • Loop k times.
  • In each iteration, extract the smallest element using poll().
  • Increment the extracted element by 1.
  • Add the new value back to the PriorityQueue.
  • After the loop, initialize a long variable product to 1.
  • While the PriorityQueue is not empty, poll() an element, multiply it with product, and take the modulo 10^9 + 7.
  • Return the final product.

Walkthrough

The strategy remains the same: greedily increment the smallest element.

  1. First, we build a min-heap from all the elements in the nums array. This allows us to retrieve the minimum element in O(log N) time.
  2. We then loop k times. In each iteration, we: a. Extract the minimum element from the heap using poll(). b. Increment it by 1. c. Insert the incremented element back into the heap using add().
  3. After k increments, the heap contains the final set of numbers. We then calculate their product. We can do this by repeatedly polling from the heap until it's empty, multiplying the elements, and taking the modulo at each step.
import java.util.PriorityQueue; class Solution {    public int maximumProduct(int[] nums, int k) {        int MOD = 1_000_000_007;        PriorityQueue<Long> pq = new PriorityQueue<>();        for (int num : nums) {            pq.add((long)num);        }         for (int i = 0; i < k; i++) {            long smallest = pq.poll();            pq.add(smallest + 1);        }         long product = 1;        while (!pq.isEmpty()) {            product = (product * pq.poll()) % MOD;        }        return (int) product;    }}

Complexity

Time

O((N+k) log N). Building the heap from N elements takes O(N log N) time. Then, we perform `k` operations, each involving a `poll` and an `add`, both of which take O(log N) time. This part is O(k log N). The final product calculation involves polling N elements, taking O(N log N). The total complexity is dominated by these operations.

Space

O(N), as the `PriorityQueue` needs to store all N elements from the input array.

Trade-offs

Pros

  • Significantly more efficient than the brute-force approach.

  • Guaranteed to pass within the time limits for the given constraints.

  • Relatively easy to implement correctly.

Cons

  • Uses extra space to store the heap.

  • Can be further optimized if k is very large, as it processes increments one by one.

Solutions

class Solution {private  static final int MOD = (int)1 e9 + 7;public  int maximumProduct(int[] nums, int k) {    PriorityQueue<Integer> q = new PriorityQueue<>();    for (int v : nums) {      q.offer(v);    }    while (k-- > 0) {      q.offer(q.poll() + 1);    }    long ans = 1;    while (!q.isEmpty()) {      ans = (ans * q.poll()) % MOD;    }    return (int)ans;  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.