Question

2233. K 次增加后的最大乘积

English Version

题目描述

给你一个非负整数数组 nums 和一个整数 k 。每次操作，你可以选择 nums 中 任一 元素并将它 增加 1 。

请你返回 至多 k 次操作后，能得到的_ _nums的 最大乘积 。由于答案可能很大，请你将答案对 109 + 7 取余后返回。

 

示例 1：

输入：nums = [0,4], k = 5
输出：20
解释：将第一个数增加 5 次。
得到 nums = [5, 4] ，乘积为 5 * 4 = 20 。
可以证明 20 是能得到的最大乘积，所以我们返回 20 。
存在其他增加 nums 的方法，也能得到最大乘积。

示例 2：

输入：nums = [6,3,3,2], k = 2
输出：216
解释：将第二个数增加 1 次，将第四个数增加 1 次。
得到 nums = [6, 4, 3, 3] ，乘积为 6 * 4 * 3 * 3 = 216 。
可以证明 216 是能得到的最大乘积，所以我们返回 216 。
存在其他增加 nums 的方法，也能得到最大乘积。

 

提示：

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

解法

方法一：贪心 + 优先队列（小根堆）

每次操作，贪心地选择最小的元素进行加 $1$，共进行 $k$ 次操作。最后累乘所有元素得到结果。注意取模操作。

时间复杂度 $O(n+klogn)$。其中，$n$ 表示 $nums$ 的长度。建堆的时间复杂度为 $O(n)$，每次弹出最小元素进行加 $1$，再放回堆中，时间复杂度为 $O(logn)$，共进行 $k$ 次操作。

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

class Solution {
  private static final int MOD = (int) 1e9 + 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;
  }
}

class Solution {
public:
  int maximumProduct(vector<int>& nums, int k) {
    int mod = 1e9 + 7;
    make_heap(nums.begin(), nums.end(), greater<int>());
    while (k--) {
      pop_heap(nums.begin(), nums.end(), greater<int>());
      ++nums.back();
      push_heap(nums.begin(), nums.end(), greater<int>());
    }
    long long ans = 1;
    for (int v : nums) ans = (ans * v) % mod;
    return ans;
  }
};

func maximumProduct(nums []int, k int) int {
  h := hp{nums}
  for heap.Init(&h); k > 0; k-- {
    h.IntSlice[0]++
    heap.Fix(&h, 0)
  }
  ans := 1
  for _, v := range nums {
    ans = (ans * v) % (1e9 + 7)
  }
  return ans
}

type hp struct{ sort.IntSlice }

func (hp) Push(any)   {}
func (hp) Pop() (_ any) { return }

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var maximumProduct = function (nums, k) {
  const n = nums.length;
  let pq = new MinPriorityQueue();
  for (let i = 0; i < n; i++) {
    pq.enqueue(nums[i]);
  }
  for (let i = 0; i < k; i++) {
    pq.enqueue(pq.dequeue().element + 1);
  }
  let ans = 1;
  const limit = 10 ** 9 + 7;
  for (let i = 0; i < n; i++) {
    ans = (ans * pq.dequeue().element) % limit;
  }
  return ans;
};