3266. Final Array State After K Multiplication Operations II

• Time: $O(\texttt{sort} + k\log n + n\log(k / n))$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  vector<int> getFinalState(vector<int>& nums, int k, int multiplier) {
    if (multiplier == 1)
      return nums;

    const int n = nums.size();
    const int maxNum = ranges::max(nums);
    vector<int> ans(n);
    using P = pair<int, int>;  // (nums[i], i)
    priority_queue<P, vector<P>, greater<>> minHeap;

    for (int i = 0; i < n; ++i)
      minHeap.emplace(nums[i], i);

    // Keep multiplying the minimum number as close as possible to the maximum
    // number in the array. After that, stop multiplying the minimum number
    // because it will be greater than the maximum number in the array and break
    // the circularity.
    while (k > 0 &&
           static_cast<long>(minHeap.top().first) * multiplier <= maxNum) {
      const auto [num, i] = minHeap.top();
      minHeap.pop();
      minHeap.emplace(num * multiplier, i);
      --k;
    }

    vector<pair<int, int>> sortedIndexedNums;
    while (!minHeap.empty())
      sortedIndexedNums.push_back(minHeap.top()), minHeap.pop();

    const int multipliesPerNum = k / n;
    const int remainingK = k % n;

    // Evenly distribute the remaining multiplications to each number by using
    // fast exponentiation.
    for (auto& [num, _] : sortedIndexedNums)
      num = (num * modPow(multiplier, multipliesPerNum)) % kMod;

    // Distribute the remaining multiplications to the minimum `remainingK`
    // numbers.
    for (int i = 0; i < remainingK; ++i)
      sortedIndexedNums[i].first =
          (static_cast<long>(sortedIndexedNums[i].first) * multiplier % kMod);

    for (const auto& [num, i] : sortedIndexedNums)
      ans[i] = num;

    return ans;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % kMod;
    return modPow(x * x % kMod, n / 2);
  }
};

class Solution {
  public int[] getFinalState(int[] nums, int k, int multiplier) {
    if (multiplier == 1)
      return nums;

    final int n = nums.length;
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    int[] ans = new int[n];
    // (nums[i], i)
    Queue<int[]> minHeap = new PriorityQueue<>(
        Comparator.comparingInt((int[] a) -> a[0]).thenComparingInt((int[] a) -> a[1]));

    for (int i = 0; i < n; ++i)
      minHeap.offer(new int[] {nums[i], i});

    // Keep multiplying the minimum number as close as possible to the maximum
    // number in the array. After that, stop multiplying the minimum number
    // because it will be greater than the maximum number in the array and break
    // the circularity.
    while (k > 0 && (long) minHeap.peek()[0] * multiplier <= maxNum) {
      final int num = minHeap.peek()[0];
      final int i = minHeap.poll()[1];
      minHeap.offer(new int[] {num * multiplier, i});
      --k;
    }

    List<int[]> sortedIndexedNums = new ArrayList<>(minHeap);
    Collections.sort(sortedIndexedNums,
                     Comparator.comparingInt((int[] sortedIndexedNum) -> sortedIndexedNum[0])
                         .thenComparingInt((int[] sortedIndexedNum) -> sortedIndexedNum[1]));

    final int multipliesPerNum = k / n;
    final int remainingK = k % n;

    // Evenly distribute the remaining multiplications to each number by using
    // fast exponentiation.
    for (int[] indexedNums : sortedIndexedNums)
      indexedNums[0] = (int) ((long) indexedNums[0] * modPow(multiplier, multipliesPerNum) % MOD);

    // Distribute the remaining multiplications to the minimum `remainingK`
    // numbers.
    for (int i = 0; i < remainingK; ++i)
      sortedIndexedNums.get(i)[0] = (int) ((long) sortedIndexedNums.get(i)[0] * multiplier % MOD);

    for (int[] indexedNums : sortedIndexedNums) {
      final int num = indexedNums[0];
      final int i = indexedNums[1];
      ans[i] = num;
    }

    return ans;
  }

  private static final int MOD = 1_000_000_007;

  private long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % MOD;
    return modPow(x * x % MOD, n / 2);
  }
}

class Solution:
  def getFinalState(
      self,
      nums: list[int],
      k: int,
      multiplier: int
  ) -> list[int]:
    if multiplier == 1:
      return nums

    MOD = 1_000_000_007
    n = len(nums)
    maxNum = max(nums)
    ans = [0] * n
    minHeap = [(num, i) for i, num in enumerate(nums)]

    heapq.heapify(minHeap)

    # Keep multiplying the minimum number as close as possible to the maximum
    # number in the array. After that, stop multiplying the minimum number
    # because it will be greater than the maximum number in the array and break
    # the circularity.
    while k > 0 and minHeap[0][0] * multiplier <= maxNum:
      num, i = heapq.heappop(minHeap)
      heapq.heappush(minHeap, (num * multiplier, i))
      k -= 1

    sortedIndexedNums = sorted(minHeap)
    multipliesPerNum, remainingK = divmod(k, n)

    # Evenly distribute the remaining multiplications to each number by using
    # fast exponentiation.
    for index, (num, i) in enumerate(sortedIndexedNums):
      sortedIndexedNums[index] = (
          sortedIndexedNums[index][0] *
          pow(multiplier, multipliesPerNum, MOD) % MOD, i)

    # Distribute the remaining multiplications to the minimum `remainingK`
    # numbers.
    for index in range(remainingK):
      sortedIndexedNums[index] = (
          sortedIndexedNums[index][0] * multiplier % MOD,
          sortedIndexedNums[index][1])

    for num, i in sortedIndexedNums:
      ans[i] = num

    return ans