Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code : Final Array State After K Multiplication Operations II

const int mod = (int)1e9 + 7;

#define order first
#define idx second

class Solution {
  public:
    vector<int> getFinalState(vector<int> &nums, int k, int multiplier);
};

int power(int base, int exp) {
    if (exp == 0) {
        return 1;
    }

    int half = power(base, exp / 2);
    long long res = (1LL * half * half);
    res %= mod;
    if (exp % 2 != 0) {
        res *= base;
        res %= mod;
    }

    return res;
}

vector<int> Solution::getFinalState(vector<int> &a, int k, int m) {
    int n = a.size();
    if (m == 1) {
        return a;
    }

    vector<int> op_cnt(n, 0);

    // set stores {order, i}
    set<pair<long long, int>> st;
    for (int i = 0; i < n; i++) {
        st.insert({a[i], i});
    }

    long long left = (*st.begin()).order;
    long long right = (*st.rbegin()).order;
    while (k && left * m <= right) {
        auto now = *st.begin();
        st.erase(st.begin());

        now.order *= m;
        op_cnt[now.idx]++;
        k--;

        st.insert(now);
        left = (*st.begin()).order;
        right = (*st.rbegin()).order;
    }

    int q = k / n;
    int remain = k % n;

    for (int i = 0; i < remain; i++) {
        auto now = *st.begin();
        st.erase(st.begin());
        now.order *= m;
        st.insert(now);
        op_cnt[now.idx]++;
    }

    for (int i = 0; i < n; i++) {
        op_cnt[i] += q;
        a[i] = (1LL * a[i] * power(m, op_cnt[i])) % mod;
    }
    return a;
}