## DEV Community

Abhishek Chaudhary

Posted on

# 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`

SOLUTION:

``````import heapq

class Solution:
def maximumProduct(self, nums: List[int], k: int) -> int:
MAX = 10 ** 9 + 7
n = len(nums)
heap = []
deleted = {}
for i, num in enumerate(nums):
heapq.heappush(heap, (num, i))
for _ in range(k):
curr, j = heapq.heappop(heap)
while (curr, j) in deleted:
deleted[(curr, j)] -= 1
if deleted[(curr, j)] == 0:
del deleted[(curr, j)]
curr, j = heapq.heappop(heap)
nums[j] += 1
deleted[(curr, j)] = deleted.get((curr, j), 0) + 1
heapq.heappush(heap, (curr + 1, j))
p = 1
for num in nums:
p = (p * num) % MAX
return p
``````