Skip to content

2818. Apply Operations to Maximize Score 👍

  • Time: $O(k\log(\log k) + \texttt{sort}(n))$, where $k = \max(\texttt{nums})$
  • Space: $O(\max(\texttt{nums}) + n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class Solution {
 public:
  int maximumScore(vector<int>& nums, int k) {
    const int n = nums.size();
    const int maxNum = ranges::max(nums);
    const vector<int> minPrimeFactors = sieveEratosthenes(maxNum + 1);
    const vector<int> primeScores = getPrimeScores(nums, minPrimeFactors);
    int ans = 1;
    // left[i] := the next index on the left (if any) s.t.
    // primeScores[left[i]] >= primeScores[i]
    vector<int> left(n, -1);
    // right[i] := the next index on the right (if any) s.t.
    // primeScores[right[i]] > primeScores[i]
    vector<int> right(n, n);
    stack<int> stack;

    // Find the next indices on the left where `primeScores` are greater or
    // equal.
    for (int i = n - 1; i >= 0; --i) {
      while (!stack.empty() && primeScores[stack.top()] <= primeScores[i])
        left[stack.top()] = i, stack.pop();
      stack.push(i);
    }

    stack = std::stack<int>();

    // Find the next indices on the right where `primeScores` are greater.
    for (int i = 0; i < n; ++i) {
      while (!stack.empty() && primeScores[stack.top()] < primeScores[i])
        right[stack.top()] = i, stack.pop();
      stack.push(i);
    }

    vector<pair<int, int>> numAndIndexes;

    for (int i = 0; i < n; ++i)
      numAndIndexes.emplace_back(nums[i], i);

    ranges::sort(numAndIndexes,
                 [&](const pair<int, int>& a, const pair<int, int>& b) {
      return a.first == b.first ? a.second < b.second : a.first > b.first;
    });

    for (const auto& [num, i] : numAndIndexes) {
      // nums[i] is the maximum value in the range [left[i] + 1, right[i] - 1]
      // So, there are (i - left[i]) * (right[i] - 1) ranges where nums[i] will
      // be chosen.
      const int rangeCount = (i - left[i]) * (right[i] - i);
      const int actualCount = min(rangeCount, k);
      k -= actualCount;
      ans = (static_cast<long>(ans) * modPow(num, actualCount)) % kMod;
    }

    return ans;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }

  // Gets the minimum prime factor of i, where 1 < i <= n.
  vector<int> sieveEratosthenes(int n) {
    vector<int> minPrimeFactors(n + 1);
    iota(minPrimeFactors.begin() + 2, minPrimeFactors.end(), 2);
    for (int i = 2; i * i < n; ++i)
      if (minPrimeFactors[i] == i)  // `i` is prime.
        for (int j = i * i; j < n; j += i)
          minPrimeFactors[j] = min(minPrimeFactors[j], i);
    return minPrimeFactors;
  }

  vector<int> getPrimeScores(const vector<int>& nums,
                             const vector<int>& minPrimeFactors) {
    vector<int> primeScores;
    for (const int num : nums)
      primeScores.push_back(getPrimeScore(num, minPrimeFactors));
    return primeScores;
  }

  int getPrimeScore(int num, const vector<int>& minPrimeFactors) {
    unordered_set<int> primeFactors;
    while (num > 1) {
      const int divisor = minPrimeFactors[num];
      primeFactors.insert(divisor);
      while (num % divisor == 0)
        num /= divisor;
    }
    return primeFactors.size();
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
class Solution {
  public int maximumScore(List<Integer> nums, int k) {
    final int n = nums.size();
    final int maxNum = Collections.max(nums);
    final int[] minPrimeFactors = sieveEratosthenes(maxNum + 1);
    final int[] primeScores = getPrimeScores(nums, minPrimeFactors);
    int ans = 1;
    // left[i] := the next index on the left (if any)
    //            s.t. primeScores[left[i]] >= primeScores[i]
    int[] left = new int[n];
    Arrays.fill(left, -1);
    // right[i] := the next index on the right (if any)
    //             s.t. primeScores[right[i]] > primeScores[i]
    int[] right = new int[n];
    Arrays.fill(right, n);
    Deque<Integer> stack = new ArrayDeque<>();

    // Find the next indices on the left where `primeScores` are greater or equal.
    for (int i = n - 1; i >= 0; --i) {
      while (!stack.isEmpty() && primeScores[stack.peek()] <= primeScores[i])
        left[stack.pop()] = i;
      stack.push(i);
    }

    stack.clear();

    // Find the next indices on the right where `primeScores` are greater.
    for (int i = 0; i < n; ++i) {
      while (!stack.isEmpty() && primeScores[stack.peek()] < primeScores[i])
        right[stack.pop()] = i;
      stack.push(i);
    }

    Pair<Integer, Integer>[] numAndIndexes = new Pair[n];

    for (int i = 0; i < n; ++i)
      numAndIndexes[i] = new Pair<>(nums.get(i), i);

    Arrays.sort(numAndIndexes,
                (a, b)
                    -> a.getKey().equals(b.getKey()) ? a.getValue().compareTo(b.getValue())
                                                     : b.getKey().compareTo(a.getKey()));

    for (Pair<Integer, Integer> numAndIndex : numAndIndexes) {
      final int num = numAndIndex.getKey();
      final int i = numAndIndex.getValue();
      // nums[i] is the maximum value in the range [left[i] + 1, right[i] - 1]
      // So, there are (i - left[i]) * (right[i] - 1) ranges where nums[i] will
      // be chosen.
      final int rangeCount = (i - left[i]) * (right[i] - i);
      final int actualCount = Math.min(rangeCount, k);
      k -= actualCount;
      ans = (int) ((1L * ans * modPow(num, actualCount)) % kMod);
    }

    return ans;
  }

  private static final int kMod = 1_000_000_007;

  private long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % kMod;
    return modPow(x * x % kMod, n / 2);
  }

  // Gets the minimum prime factor of i, where 1 < i <= n.
  private int[] sieveEratosthenes(int n) {
    int[] minPrimeFactors = new int[n + 1];
    for (int i = 2; i <= n; ++i)
      minPrimeFactors[i] = i;
    for (int i = 2; i * i < n; ++i)
      if (minPrimeFactors[i] == i) // `i` is prime.
        for (int j = i * i; j < n; j += i)
          minPrimeFactors[j] = Math.min(minPrimeFactors[j], i);
    return minPrimeFactors;
  }

  private int[] getPrimeScores(List<Integer> nums, int[] minPrimeFactors) {
    int[] primeScores = new int[nums.size()];
    for (int i = 0; i < nums.size(); ++i)
      primeScores[i] = getPrimeScore(nums.get(i), minPrimeFactors);
    return primeScores;
  }

  private int getPrimeScore(int num, int[] minPrimeFactors) {
    Set<Integer> primeFactors = new HashSet<>();
    while (num > 1) {
      final int divisor = minPrimeFactors[num];
      primeFactors.add(divisor);
      while (num % divisor == 0)
        num /= divisor;
    }
    return primeFactors.size();
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Solution:
  def maximumScore(self, nums: list[int], k: int) -> int:
    kMod = 1_000_000_007
    n = len(nums)
    ans = 1
    minPrimeFactors = self._sieveEratosthenes(max(nums) + 1)
    primeScores = [self._getPrimeScore(num, minPrimeFactors) for num in nums]
    # left[i] := the next index on the left (if any)
    #            s.t. primeScores[left[i]] >= primeScores[i]
    left = [-1] * n
    # right[i] := the next index on the right (if any)
    #             s.t. primeScores[right[i]] > primeScores[i]
    right = [n] * n
    stack = []

    # Find the next indices on the left where `primeScores` are greater or equal.
    for i in reversed(range(n)):
      while stack and primeScores[stack[-1]] <= primeScores[i]:
        left[stack.pop()] = i
      stack.append(i)

    stack = []

    # Find the next indices on the right where `primeScores` are greater.
    for i in range(n):
      while stack and primeScores[stack[-1]] < primeScores[i]:
        right[stack.pop()] = i
      stack.append(i)

    numAndIndexes = [(num, i) for i, num in enumerate(nums)]

    def modPow(x: int, n: int) -> int:
      if n == 0:
        return 1
      if n % 2 == 1:
        return x * modPow(x, n - 1) % kMod
      return modPow(x * x % kMod, n // 2)

    for num, i in sorted(numAndIndexes, key=lambda x: (-x[0], x[1])):
      # nums[i] is the maximum value in the range [left[i] + 1, right[i] - 1]
      # So, there are (i - left[i]) * (right[i] - 1) ranges where nums[i] will
      # be chosen.
      rangeCount = (i - left[i]) * (right[i] - i)
      actualCount = min(rangeCount, k)
      k -= actualCount
      ans *= modPow(num, actualCount)
      ans %= kMod

    return ans

  def _sieveEratosthenes(self, n: int) -> list[int]:
    """Gets the minimum prime factor of i, where 2 <= i <= n."""
    minPrimeFactors = [i for i in range(n + 1)]
    for i in range(2, int(n**0.5) + 1):
      if minPrimeFactors[i] == i:  # `i` is prime.
        for j in range(i * i, n, i):
          minPrimeFactors[j] = min(minPrimeFactors[j], i)
    return minPrimeFactors

  def _getPrimeScore(self, num: int, minPrimeFactors: list[int]) -> int:
    primeFactors = set()
    while num > 1:
      divisor = minPrimeFactors[num]
      primeFactors.add(divisor)
      while num % divisor == 0:
        num //= divisor
    return len(primeFactors)