Skip to main content
Tags:

2233. Maximum Product After K Increments

link

Method 1

Sort and presum

class Solution:
    def maximumProduct(self, nums: List[int], k: int) -> int:
        
        n = len(nums)

        nums.sort()

        presum = [0] * n
        presum[0] = nums[0]

        for i in range(1, n):
            presum[i] = presum[i-1] + nums[i]

        diff = [0] * n

        for i in range(1, n):
            diff[i] = nums[i] * (i + 1) - presum[i]

        idx = bisect.bisect_left(diff, k) - 1

        total = presum[idx] + k
        num = total // (idx + 1)
        extra = total - num * (idx+1)

        res = 1

        for i in range(extra):
            res *= (num + 1)
            res %= (10**9+7)
        
        for i in range(idx+1-extra):
            res *= num
            res %= (10**9+7)

        for i in range(idx+1, n):
            res *= nums[i]
            res %= (10**9+7)

        return res % (10**9+7)


Method 2

heap

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