Apply Operations to Maximize Score

Problem

You are given an array nums of n positive integers and an integer k.

Initially, you start with a score of 1. You have to maximize your score by applying the following operation at most k times:

• Choose any non-empty subarray nums[l, ..., r] that you haven't chosen previously.
• Choose an element x of nums[l, ..., r] with the highest prime score. If multiple such elements exist, choose the one with the smallest index.
• Multiply your score by x.

Here, nums[l, ..., r] denotes the subarray of nums starting at index l and ending at the index r, both ends being inclusive.

The prime score of an integer x is equal to the number of distinct prime factors of x. For example, the prime score of 300 is 3 since 300 = 2 * 2 * 3 * 5 * 5.

Return the maximum possible score after applying at most k operations.

Since the answer may be large, return it modulo 10^9 + 7.

Examples

Example 1:

Input: nums = [8,3,9,3,8], k = 2
Output: 81
Explanation: To get a score of 81, we can apply the following operations:
- Choose subarray nums[2, ..., 2]. nums[2] is the only element in this subarray. Hence, we multiply the score by nums[2]. The score becomes 1 * 9 = 9.
- Choose subarray nums[2, ..., 3]. Both nums[2] and nums[3] have a prime score of 1, but nums[2] has the smaller index. Hence, we multiply the score by nums[2]. The score becomes 9 * 9 = 81.
It can be proven that 81 is the highest score one can obtain.

Example 2:

Input: nums = [19,12,14,6,10,18], k = 3
Output: 4788
Explanation: To get a score of 4788, we can apply the following operations: 
- Choose subarray nums[0, ..., 0]. nums[0] is the only element in this subarray. Hence, we multiply the score by nums[0]. The score becomes 1 * 19 = 19.
- Choose subarray nums[5, ..., 5]. nums[5] is the only element in this subarray. Hence, we multiply the score by nums[5]. The score becomes 19 * 18 = 342.
- Choose subarray nums[2, ..., 3]. Both nums[2] and nums[3] have a prime score of 2, but nums[2] has the smaller index. Hence, we multipy the score by nums[2]. The score becomes 342 * 14 = 4788.
It can be proven that 4788 is the highest score one can obtain.

Constraints:

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

Solution

Method 1 - Using Sieve of Eratosthenes and Sorting

The goal is to maximise the score by selecting elements from subarrays based on their prime scores. The primary challenges include:

1. Prime Score Calculation:
   • Each number's prime score is defined by the count of distinct prime factors.
   • Efficient computation of prime factors is critical for large inputs.
2. Subarray Dominance:
   • For any chosen element, identify the subarrays where it is "dominant" (having the highest prime score in the subarray).
   • Use monotonic stack techniques to determine dominant subarray regions efficiently.
3. Optimising Operations:
   • Sort elements by value in descending order and process elements to maximise score while applying up to k operations.
   • Raise the chosen element to the power of how many subarrays it dominates, modulo (10^9 + 7).

Approach

1. Efficient Prime Score Calculation:
   • Use the Sieve of Eratosthenes to generate all prime numbers up to the maximum value in the array.
   • For each number, iterate through primes to count distinct prime factors efficiently.
2. Finding Dominant Subarrays:
   • Use monotonic stack to compute:
     • Next Dominant Index: The nearest index to the right where the prime score exceeds the current element's prime score.
     • Previous Dominant Index: The nearest index to the left where the prime score exceeds the current element's prime score.
   • Using these indices, calculate the count of subarrays where each element is dominant.
3. Maximising the Score:
   • Sort elements by value (in descending order).
   • Process each element (largest first) and calculate its contribution using the number of subarrays it dominates and the remaining operations (k).

Code

Java

class Solution {

    private static final int MOD = 1_000_000_007;

    public int maximumScore(List<Integer> nums, int k) {
        int n = nums.size();
        int[] primeScores = new int[n];

        // Find the maximum element in nums to determine the range for prime generation
        int maxElement = Collections.max(nums);

        // Get all prime numbers up to maxElement using the Sieve of Eratosthenes
        List<Integer> primes = getPrimes(maxElement);

        // Calculate the prime score for each number in nums
        for (int index = 0; index < n; index++) {
            int num = nums.get(index);

            // Iterate over the generated primes to count unique prime factors
            for (int prime : primes) {
                if (prime * prime > num) break; // Stop early if prime^2 exceeds num
                if (num % prime != 0) continue; // Skip if the prime is not a factor

                primeScores[index]++; // Increment prime score for the factor
                while (num % prime == 0) num /= prime; // Remove all occurrences of this factor
            }

            // If num is still greater than 1, it is a prime number itself
            if (num > 1) primeScores[index]++;
        }

        // Initialize next and previous dominant index arrays
        int[] nextDominant = new int[n];
        int[] prevDominant = new int[n];
        Arrays.fill(nextDominant, n);
        Arrays.fill(prevDominant, -1);

        // Stack to store indices for a monotonic decreasing prime score
        Stack<Integer> decreasingPrimeScoreStack = new Stack<>();

        // Calculate the next and previous dominant indices for each number
        for (int index = 0; index < n; index++) {
            while (
                !decreasingPrimeScoreStack.isEmpty() &&
                primeScores[decreasingPrimeScoreStack.peek()] <
                primeScores[index]
            ) {
                int topIndex = decreasingPrimeScoreStack.pop();
                nextDominant[topIndex] = index;
            }

            if (!decreasingPrimeScoreStack.isEmpty()) {
                prevDominant[index] = decreasingPrimeScoreStack.peek();
            }

            decreasingPrimeScoreStack.push(index);
        }

        // Calculate the number of subarrays in which each element is dominant
        long[] numOfSubarrays = new long[n];
        for (int index = 0; index < n; index++) {
            numOfSubarrays[index] =
                (long) (nextDominant[index] - index) *
                (index - prevDominant[index]);
        }

        // Sort elements in decreasing order based on their values
        List<int[]> sortedArray = new ArrayList<>();
        for (int index = 0; index < n; index++) {
            sortedArray.add(new int[] { nums.get(index), index });
        }
        sortedArray.sort((a, b) -> Integer.compare(b[0], a[0])); // Sort in descending order

        long score = 1;
        int processingIndex = 0;

        // Process elements while there are operations left
        while (k > 0) {
            int[] element = sortedArray.get(processingIndex++);
            int num = element[0];
            int index = element[1];

            // Calculate the number of operations to apply on the current element
            long operations = Math.min(k, numOfSubarrays[index]);

            // Update the score by raising the element to the power of operations
            score = (score * power(num, operations)) % MOD;

            // Reduce the remaining operations count
            k -= operations;
        }

        return (int) score;
    }

    // Helper function to compute the power of a number modulo MOD
    private long power(long base, long exponent) {
        long res = 1;

        while (exponent > 0) {
            if (exponent % 2 == 1) {
                res = (res * base) % MOD;
            }
            base = (base * base) % MOD;
            exponent /= 2;
        }

        return res;
    }

    // Function to generate prime numbers up to a given limit using the Sieve of Eratosthenes
    private List<Integer> getPrimes(int limit) {
        boolean[] isPrime = new boolean[limit + 1];
        Arrays.fill(isPrime, true);
        List<Integer> primes = new ArrayList<>();

        for (int number = 2; number <= limit; number++) {
            if (!isPrime[number]) continue;

            primes.add(number);

            for (
                long multiple = (long) number * number;
                multiple <= limit;
                multiple += number
            ) {
                isPrime[(int) multiple] = false;
            }
        }

        return primes;
    }
}

Python

MOD = 10**9 + 7

class Solution:
    def maximumScore(self, nums: List[int], k: int) -> int:
        def get_primes(limit: int) -> List[int]:
            """Generate all prime numbers up to a given limit using the Sieve of Eratosthenes."""
            is_prime = [True] * (limit + 1)
            primes = []
            for num in range(2, limit + 1):
                if is_prime[num]:
                    primes.append(num)
                    for multiple in range(num * num, limit + 1, num):
                        is_prime[multiple] = False
            return primes

        def count_prime_score(num: int, primes: List[int]) -> int:
            """Count the distinct prime factors of a number."""
            score = 0
            for prime in primes:
                if prime * prime > num: break
                if num % prime == 0:
                    score += 1
                    while num % prime == 0:
                        num //= prime
            if num > 1:  # Leftover prime number
                score += 1
            return score
        
        def calculate_power(base: int, exponent: int) -> int:
            """Fast modular exponentiation to calculate (base^exponent) % MOD."""
            result = 1
            while exponent > 0:
                if exponent % 2 == 1:
                    result = (result * base) % MOD
                base = (base * base) % MOD
                exponent //= 2
            return result
        
        # Step 1: Calculate prime scores for all elements in nums
        max_val = max(nums)
        primes = get_primes(max_val)
        n = len(nums)
        prime_scores = [count_prime_score(num, primes) for num in nums]
        
        # Step 2: Determine next and previous dominant indices using monotonic stack
        next_dominant = [n] * n
        prev_dominant = [-1] * n
        stack = []
        
        for i in range(n):
            while stack and prime_scores[stack[-1]] < prime_scores[i]:
                next_dominant[stack.pop()] = i
            if stack:
                prev_dominant[i] = stack[-1]
            stack.append(i)
        
        # Step 3: Calculate the number of subarrays where each element is dominant
        subarray_count = [
            (next_dominant[i] - i) * (i - prev_dominant[i])
            for i in range(n)
        ]
        
        # Step 4: Sort elements by value and process them to maximise score
        sorted_elements = sorted(
            [(nums[i], i) for i in range(n)], 
            key=lambda x: -x[0]
        )
        
        score = 1
        index = 0
        
        while k > 0:
            value, idx = sorted_elements[index]
            count = min(k, subarray_count[idx])
            score = (score * calculate_power(value, count)) % MOD
            k -= count
            index += 1
        
        return score

Complexity

• ⏰ Time complexity: O(maxElement * log(log(maxElement)) + n * sqrt(maxElement) + n log n + k).
  • Prime Calculation:
    • Generating primes with the sieve: O(maxElement * log(log(maxElement))).
    • Calculating prime scores for all elements: O(n * sqrt(maxElement)).
  • Dominant Indices:
    • Using stacks: O(n).
  • Sorting:
    • Sorting elements: O(n log n).
  • Processing Elements:
    • Processing up to k elements: O(k).
• 🧺 Space complexity: O(maxElement / log(maxElement) + n).
  • Primes storage: O(maxElement / log(maxElement)).
  • Arrays for indices, scores, and subarray counts: O(n).