Apply Operations to Maximize Score — Hard Problem & Solution

The prime score of an integer is its number of distinct prime factors — the prime score of 300 is 3, since 300 is 2² · 3 · 5². Your score starts at 1.

Difficulty: Hard
Topics: Arrays, Math, Greedy, Sorting, Stack, Number Theory, Monotonic Stack
Asked at: Amazon, Google, Microsoft
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

The prime score of an integer is its number of distinct prime factors — the prime score of 300 is 3, since 300 is 2² · 3 · 5².

Your score starts at 1. You may apply the following operation at most k times:

Pick a subarray nums[l..r] that you have not picked before.
Within it, choose the element with the highest prime score; on a tie, choose the one with the smallest index.
Multiply your score by that element.

Return the maximum score, modulo 10⁹ + 7.

Example 1

Input: nums = [8,3,9,3,8], k = 2
Output: 81
Explanation: Every value has prime score 1, and `9` wins three different subarrays — so it can be taken twice.

Example 2

Input: nums = [19,12,14,6,10,18], k = 3
Output: 4788
Explanation: `19` has prime score 1 so it wins only `[19]`; `18`, then `14`, are the next best. 19 · 18 · 14 = 4788.

Example 3

Input: nums = [2,3,5], k = 1
Output: 5

Constraints

1 <= nums.length <= 10^5
1 <= nums[i] <= 10^5
1 <= k <= min(n * (n + 1) / 2, 10^9)

How to solve Apply Operations to Maximize Score

Separate the two halves of the problem. First count, for each index, how many subarrays would select it — a classic "contribution of each element as the window maximum" computed with two monotonic stacks. Then the operations are unconstrained: sort the values descending and spend k on the largest, up to each one's count, accumulating the product with fast exponentiation.

Approach

Sieve the distinct-prime-factor count up to max(nums): for each prime p, add one to every multiple of p.
With a monotonic stack left to right, find left[i] — the nearest index before i whose score is ≥ s[i] (the tie-break sends ties left).
With a stack right to left, find right[i] — the nearest index after i whose score is strictly > s[i].
count[i] = (i - left[i]) * (right[i] - i).
Sort indices by nums[i] descending; for each, take min(k, count[i]) copies via powmod, multiply into the answer and reduce k.
Stop as soon as k reaches 0.

Why it works

The asymmetry between ≥ on the left and > on the right is exactly the leftmost-on-tie rule: an equal score to the left would beat i, so it blocks; an equal score to the right would lose, so it does not. Getting this backwards double-counts subarrays whenever two equal scores sit side by side. Once the counts are known the greedy is safe, because every operation is independent — nothing you take restricts what you can take next, so the largest available value is always the right choice.

Complexity

Time — O(n log n + M log log M) where M is max(nums)
Space — O(n + M)

Pitfalls

count[i] can reach about 5 · 10⁹, so it must be held in a 64-bit integer even though the answer fits an int after the modulo.
Multiplying two residues near 10⁹ overflows 53-bit floats — split the multiplication or use 64-bit integers.
Take the product modulo 10⁹ + 7 throughout; the running score is astronomically large otherwise.
Prime score counts distinct primes: 8 = 2³ scores 1, not 3.

Reference solution

Python

from typing import List

def maximumScore(nums: List[int], k: int) -> int:
    MOD = 10**9 + 7
    n = len(nums)
    maxv = max(nums)
    sieve = [0] * (maxv + 1)
    for p in range(2, maxv + 1):
        if sieve[p] == 0:
            for q in range(p, maxv + 1, p):
                sieve[q] += 1
    s = [sieve[v] for v in nums]
    left = [-1] * n
    right = [n] * n
    stack = []
    for i in range(n):
        while stack and s[stack[-1]] < s[i]:
            stack.pop()
        left[i] = stack[-1] if stack else -1
        stack.append(i)
    stack = []
    for i in range(n - 1, -1, -1):
        while stack and s[stack[-1]] <= s[i]:
            stack.pop()
        right[i] = stack[-1] if stack else n
        stack.append(i)
    order = sorted(range(n), key=lambda i: -nums[i])
    ans = 1
    for i in order:
        if k == 0:
            break
        count = (i - left[i]) * (right[i] - i)
        take = min(count, k)
        ans = ans * pow(nums[i], take, MOD) % MOD
        k -= take
    return ans

JavaScript

var maximumScore = function(nums, k) {
    var MOD = 1000000007;
    var mulmod = function(a, b) {
        var ah = Math.floor(a / 65536), al = a % 65536;
        return ((ah * b % MOD) * 65536 + al * b) % MOD;
    };
    var powmod = function(base, exp) {
        var result = 1, bb = base % MOD, e = exp;
        while (e > 0) {
            if (e % 2 === 1) result = mulmod(result, bb);
            bb = mulmod(bb, bb);
            e = Math.floor(e / 2);
        }
        return result;
    };
    var n = nums.length, i, maxv = 0;
    for (i = 0; i < n; i++) if (nums[i] > maxv) maxv = nums[i];
    var sieve = [];
    for (i = 0; i <= maxv; i++) sieve.push(0);
    for (var p = 2; p <= maxv; p++) {
        if (sieve[p] === 0) for (var q = p; q <= maxv; q += p) sieve[q]++;
    }
    var s = [], left = [], right = [];
    for (i = 0; i < n; i++) { s.push(sieve[nums[i]]); left.push(-1); right.push(n); }
    var stack = [];
    for (i = 0; i < n; i++) {
        while (stack.length > 0 && s[stack[stack.length - 1]] < s[i]) stack.pop();
        left[i] = stack.length > 0 ? stack[stack.length - 1] : -1;
        stack.push(i);
    }
    stack = [];
    for (i = n - 1; i >= 0; i--) {
        while (stack.length > 0 && s[stack[stack.length - 1]] <= s[i]) stack.pop();
        right[i] = stack.length > 0 ? stack[stack.length - 1] : n;
        stack.push(i);
    }
    var order = [];
    for (i = 0; i < n; i++) order.push(i);
    order.sort(function(a, b) { return nums[b] - nums[a]; });
    var remaining = k, ans = 1;
    for (var t = 0; t < n && remaining > 0; t++) {
        var idx = order[t];
        var count = (idx - left[idx]) * (right[idx] - idx);
        var take = count < remaining ? count : remaining;
        ans = mulmod(ans, powmod(nums[idx], take));
        remaining -= take;
    }
    return ans;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue