力扣 3266. K 次乘运算后的最终数组 II

发现 \(k\le10^9\)，肯定不能直接模拟了。但是注意到，如果最小值 \(\min\times m>\max\) 存在，我们可以直接给每个数乘上 \(m^{\lfloor\frac{k}{n}\rfloor}\)，最后再处理剩下的 \(k\bmod n\) 个 \(m\) 即可，这个复杂度是 \(O(n)\) 的。具体地说，

• 最小值乘上一个 \(m\) 超过最大值之后，其他的数也可以乘一个 \(m\) 超过最大值。
• 因为 \(\min<\max\)，所以最大值只乘一个 \(m\) 又能成为最大值。
• 因为 \(\min\times m>\max\)，所以 \(\min\times m^2>\max\times m\)，最小值又可以通过乘一个 \(m\) 超过最大值……

这个过程可以继续进行下去，直到 \(k\) 不够用。

那么什么时候 \(\min\times m>\max\) 成立呢？发现这个条件随着不断模拟 \(\min\times m\) 是一定能达到的，而且步数并不多。\(nums_i\le 10^9\)，\(m\) 取最小值 \(2\)，也只需要乘 \(31\) 次就能达到；最多 \(n-1\) 个数每个乘 \(31\) 次，也不会超时。

所以一开始直接小根堆模拟，达到 \(\min\times m>\max\) 后用快速幂批量处理，最后剩下的部分继续模拟完即可。

using LL = long long;
const int MOD = 1'000'000'007;
LL power(LL a, LL n); // 快速幂（略）

vector<int> getFinalState(vector<int>& nums, int k, int m) {
    if (m == 1) {
        return move(nums);
    }
    int n = nums.size();
    int mx = ranges::max(nums);
    vector<pair<LL, int>> v(n);
    for (int i = 0; i < n; i++) {
        v[i] = {nums[i], i};
    }
    ranges::make_heap(v, greater{}); // 小根堆

    for (; k && v[0].first < mx; k--) {
        (ranges::pop_heap(v, greater{}) - 1)->first *= m; // 取堆头，放到尾部
        ranges::push_heap(v, greater{}); // 最后一个元素入堆（一般配合push_back）
    }

    LL p1 = power(m, k / n); // 后面的乘 m^(k/n)
    LL p2 = p1 * m % MOD; // 前 k%n 个乘 m^(k/n+1)
    ranges::nth_element(v, v.begin() + k % n); // 前 k 小放到前面，比 sort 快
    for (int i = 0; i < n; i++) {
        auto& [x, j] = v[i];
        nums[j] = x % MOD * (i < k % n ? p2 : p1) % MOD;
    }
    return move(nums);
}