Apply Operations to Maximize Score

hard monotonic stack greedy number theory

Problem

Start with a score of 1. Up to k times, pick a previously unused subarray nums[l..r], find the element with the highest prime score (number of distinct prime factors; ties broken by smallest index), and multiply the score by that element. The prime score of 300 = 2·2·3·5·5 is 3. Return the maximum score modulo 109+7.

Each index i is the “winner” of every subarray where its prime score dominates. A monotonic stack gives the span (left[i], right[i]) over which it wins, so it can be chosen ranges[i] = (i−left[i])·(right[i]−i) times. Spend operations greedily on the largest values first.

Inputnums = [8,3,9,3,8], k = 2
Output81
All prime scores are 1. Index 2 (value 9) wins 3 ranges; spend both ops there: 9·9 = 81.

def maximumScore(nums, k):
    MOD = 10**9 + 7
    n = len(nums)

    def prime_score(x):                 # distinct prime factors
        c, d = 0, 2
        while d * d <= x:
            if x % d == 0:
                c += 1
                while x % d == 0: x //= d
            d += 1
        return c + (1 if x > 1 else 0)

    s = [prime_score(v) for v in nums]
    left, right = [-1] * n, [n] * n
    st = []                             # left: nearest s >= s[i]
    for i in range(n):
        while st and s[st[-1]] < s[i]: st.pop()
        left[i] = st[-1] if st else -1
        st.append(i)
    st = []                             # right: nearest s > s[i]
    for i in range(n - 1, -1, -1):
        while st and s[st[-1]] <= s[i]: st.pop()
        right[i] = st[-1] if st else n
        st.append(i)
    rng = [(i - left[i]) * (right[i] - i) for i in range(n)]

    score, rem = 1, k                   # greedily spend on big values
    for i in sorted(range(n), key=lambda i: -nums[i]):
        if rem <= 0: break
        take = min(rng[i], rem)
        score = score * pow(nums[i], take, MOD) % MOD
        rem -= take
    return score
function maximumScore(nums, k) {
  const MOD = 1000000007n, n = nums.length;
  const primeScore = (x) => {            // distinct prime factors
    let c = 0;
    for (let d = 2; d * d <= x; d++)
      if (x % d === 0) { c++; while (x % d === 0) x /= d; }
    return c + (x > 1 ? 1 : 0);
  };
  const s = nums.map(primeScore);
  const left = Array(n).fill(-1), right = Array(n).fill(n);
  let st = [];                          // left: nearest s >= s[i]
  for (let i = 0; i < n; i++) {
    while (st.length && s[st[st.length - 1]] < s[i]) st.pop();
    left[i] = st.length ? st[st.length - 1] : -1;
    st.push(i);
  }
  st = [];                              // right: nearest s > s[i]
  for (let i = n - 1; i >= 0; i--) {
    while (st.length && s[st[st.length - 1]] <= s[i]) st.pop();
    right[i] = st.length ? st[st.length - 1] : n;
    st.push(i);
  }
  const rng = nums.map((_, i) => (i - left[i]) * (right[i] - i));
  const order = [...nums.keys()].sort((a, b) => nums[b] - nums[a]);
  let score = 1n, rem = k;              // greedily spend on big values
  const power = (b, e) => {
    let r = 1n; b %= MOD;
    while (e > 0n) { if (e & 1n) r = r * b % MOD; b = b * b % MOD; e >>= 1n; }
    return r;
  };
  for (const i of order) {
    if (rem <= 0) break;
    const take = Math.min(rng[i], rem);
    score = score * power(BigInt(nums[i]), BigInt(take)) % MOD;
    rem -= take;
  }
  return Number(score);
}
int maximumScore(int[] nums, int k) {
    final long MOD = 1_000_000_007L;
    int n = nums.length;
    int[] s = new int[n];
    for (int i = 0; i < n; i++) s[i] = primeScore(nums[i]);
    int[] left = new int[n], right = new int[n];
    Deque<Integer> st = new ArrayDeque<>();   // left: nearest s >= s[i]
    for (int i = 0; i < n; i++) {
        while (!st.isEmpty() && s[st.peek()] < s[i]) st.pop();
        left[i] = st.isEmpty() ? -1 : st.peek();
        st.push(i);
    }
    st.clear();                              // right: nearest s > s[i]
    for (int i = n - 1; i >= 0; i--) {
        while (!st.isEmpty() && s[st.peek()] <= s[i]) st.pop();
        right[i] = st.isEmpty() ? n : st.peek();
        st.push(i);
    }
    long[] rng = new long[n];
    for (int i = 0; i < n; i++) rng[i] = (long)(i - left[i]) * (right[i] - i);
    Integer[] ord = new Integer[n];
    for (int i = 0; i < n; i++) ord[i] = i;
    Arrays.sort(ord, (a, b) -> nums[b] - nums[a]);
    long score = 1, rem = k;                 // greedily spend on big values
    for (int i : ord) {
        if (rem <= 0) break;
        long take = Math.min(rng[i], rem);
        score = score * power(nums[i], take, MOD) % MOD;
        rem -= take;
    }
    return (int) score;
}
int maximumScore(vector<int>& nums, int k) {
    const long long MOD = 1e9 + 7;
    int n = nums.size();
    vector<int> s(n);
    for (int i = 0; i < n; i++) s[i] = primeScore(nums[i]);
    vector<int> left(n, -1), right(n, n), st;
    for (int i = 0; i < n; i++) {             // left: nearest s >= s[i]
        while (!st.empty() && s[st.back()] < s[i]) st.pop_back();
        left[i] = st.empty() ? -1 : st.back();
        st.push_back(i);
    }
    st.clear();                               // right: nearest s > s[i]
    for (int i = n - 1; i >= 0; i--) {
        while (!st.empty() && s[st.back()] <= s[i]) st.pop_back();
        right[i] = st.empty() ? n : st.back();
        st.push_back(i);
    }
    vector<long long> rng(n);
    for (int i = 0; i < n; i++) rng[i] = (long long)(i - left[i]) * (right[i] - i);
    vector<int> ord(n);
    iota(ord.begin(), ord.end(), 0);
    sort(ord.begin(), ord.end(), [&](int a, int b){ return nums[a] > nums[b]; });
    long long score = 1, rem = k;             // greedily spend on big values
    for (int i : ord) {
        if (rem <= 0) break;
        long long take = min(rng[i], rem);
        score = score * power(nums[i], take, MOD) % MOD;
        rem -= take;
    }
    return (int) score;
}
Time: O(n·√max + n log n) Space: O(n)