Problem Statement in English

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

You can choose a non empty subarray that you didn’t choose previously. You must then pick the number from that array with the highest “prime score”. If there are multiple numbers with the same prime score, you must choose the one with the smallest index.

The prime score of a number is the number of distinct prime factors of that number.

You must now multiply your answer by the number you just chose.

You can repeat this process atmost k times.

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

Approach

Whew. There’s a lot going on here. Let’s break it down.

  1. Calculating the prime score of each number: We can use a modified Sieve of Eratosthenes to find the number of distinct prime factors for each number in nums. This will allow us to determine the prime score of each number efficiently.

    The sieve works by iterating through each number, and for each number that is prime, we mark its multiples. That’s the regular sieve. The modification we make, is that for each multiple, we increment the count of distinct prime factors. This way, we can find the prime score for all numbers up to the maximum number in nums.

  2. Finding the left and right ranges: We need to find the nearest previous and next indices with strictly greater prime scores. This will help us calculate how many times each number can dominate others. This is done using a montonic stack approach.

    In order to find the left range, we iterate through the array (from left to right) and maintain a stack of indices. For each index, we pop from the stack until we find an index with a greater prime score. The left range for the index we’re processing is the first index whose prime score is greater than this one’s.

    In order to find the right range, we do the same but iterate from right to left. The right range for an index is the first index we encoutner whilst popping, whose prime score is greater than or equal to this one’s.

  3. Calculating the dominance count: For each number, we can calculate how many times it can dominate others by using the left and right ranges. The dominance count for an index i is given by the formula (right_range[i] - i) * (i - left_range[i]). This gives us the number of indices that this number can dominate. Personally, I have no idea who came up with this formula, but it works. Some things are just meant to be.

  4. Finding out the top k subarrays to use: We can use a max heap to keep track of the numbers and their dominance counts. This allows us to efficiently get the maximum product by popping the largest number from the heap and multiplying it with the result.

And with that, we’re finally done.

Solution in Python


from heapq import heapify, heappop
from collections import defaultdict
from math import ceil, sqrt
from typing import List

class Solution:
    def maximumScore(self, nums: List[int], k: int) -> int:
        MOD = 1_000_000_007

        # Edge case: empty nums
        if not nums:
            return 0

        # Find distinct prime factors count using modified Sieve of Eratosthenes
        largest = max(nums)
        factors = [0] * (largest + 1)

        for i in range(2, sqrt(largest) + 1):
            if factors[i] == 0:  # Prime number
                for j in range(i, largest + 1, i):
                    factors[j] += 1

        # Compute left and right ranges for dominance calculation
        LEN = len(nums)
        left_range, right_range = [-1] * LEN, [LEN] * LEN
        stack = []

        # Calculate left_range: nearest previous index with strictly greater factors
        for i in range(LEN):
            while stack and factors[nums[stack[-1]]] < factors[nums[i]]:
                stack.pop()
            left_range[i] = stack[-1] if stack else -1
            stack.append(i)

        stack = []

        # Calculate right_range: nearest next index with greater or equal factors
        for i in range(LEN - 1, -1, -1):
            while stack and factors[nums[stack[-1]]] <= factors[nums[i]]:
                stack.pop()
            right_range[i] = stack[-1] if stack else LEN
            stack.append(i)

        # Compute dominance count
        dominatesFor = [(right_range[i] - i) * (i - left_range[i]) for i in range(LEN)]

        # Use max heap to get the maximum product
        maxheap = [(-nums[i], dominatesFor[i]) for i in range(LEN)]
        heapify(maxheap)

        res = 1
        while k and maxheap:
            num, times = heappop(maxheap)
            take = min(k, times)
            k -= take
            res = (res * pow(-num, take, MOD)) % MOD

        return res

Complexity

  • Time: $O(klog(n))$
    This is because we are using a max heap to keep track of the numbers and their dominance counts. Each pop operation takes $O(log(n))$ time, and we do this at most k times.

  • Space: $O(n)$
    Since we are storing the prime factors and the left and right ranges in separate arrays, the space complexity is linear with respect to the size of the input array.

And we are done.