2233. Maximum Product After K Increments

2233. Maximum Product After K Increments

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

难度 : Medium

Solution

Use PriorityQueue. Put all the nums to the PriorityQueue.
For each operation, take out the min value x, and put (x+1) back to the queue.
Proof:

For example, there are 2 nums x and y
if increase x, then the product is 
(x + 1) * y = x * y + y
if increase y, then the product is 
(y + 1) * x = x * y + x

if x < y then (x*y) + y > (x*y) +x

Time Complexity : O(K logN + N logN)
Space Complexity: O(N)

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