# Maximize Count of Distinct Primes After Split
**Difficulty:** HARD
[External](https://leetcode.com/problems/maximize-count-of-distinct-primes-after-split)
Canonical: https://scaleengineer.com/dsa/problems/maximize-count-of-distinct-primes-after-split
**Patterns:** [Math](https://scaleengineer.com/dsa/patterns/math), [Number Theory](https://scaleengineer.com/dsa/patterns/number-theory)
**Data structures:** Array, Segment Tree
---
## Problem
You are given an integer array `nums` having length `n` and a 2D integer array `queries` where `queries[i] = [idx, val]`.

For each query:

1. Update `nums[idx] = val`.
2. Choose an integer `k` with `1 <= k < n` to split the array into the non-empty prefix `nums[0..k-1]` and suffix `nums[k..n-1]` such that the sum of the counts of **distinct** prime values in each part is **maximum**.

**Note:** The changes made to the array in one query persist into the next query.

Return an array containing the result for each query, in the order they are given.

**Example 1:**

**Input:** nums = \[2,1,3,1,2\], queries = \[\[1,2\],\[3,3\]\]

**Output:** \[3,4\]

**Explanation:**

* Initially `nums = [2, 1, 3, 1, 2]`.
* After 1st query, `nums = [2, 2, 3, 1, 2]`. Split `nums` into `[2]` and `[2, 3, 1, 2]`. `[2]` consists of 1 distinct prime and `[2, 3, 1, 2]` consists of 2 distinct primes. Hence, the answer for this query is `1 + 2 = 3`.
* After 2nd query, `nums = [2, 2, 3, 3, 2]`. Split `nums` into `[2, 2, 3]` and `[3, 2]` with an answer of `2 + 2 = 4`.
* The output is `[3, 4]`.

**Example 2:**

**Input:** nums = \[2,1,4\], queries = \[\[0,1\]\]

**Output:** \[0\]

**Explanation:**

* Initially `nums = [2, 1, 4]`.
* After 1st query, `nums = [1, 1, 4]`. There are no prime numbers in `nums`, hence the answer for this query is 0.
* The output is `[0]`.

**Constraints:**

* `2 <= n == nums.length <= 5 * 104`
* `1 <= queries.length <= 5 * 104`
* `1 <= nums[i] <= 105`
* `0 <= queries[i][0] < nums.length`
* `1 <= queries[i][1] <= 105`

# Approaches
## Brute Force Simulation
This is a straightforward, brute-force approach that directly simulates the process described in the problem. For each query, after updating the array, it iterates through every possible split point `k`. For each split, it computes the distinct prime counts for the resulting prefix and suffix subarrays from scratch and sums them up. The maximum sum found across all splits is the answer for that query.
**Time:** O(Q * N^2 * log(V)), where `Q` is the number of queries, `N` is the length of `nums`, and `V` is the maximum value. For each of the `Q` queries, we iterate `N` times for the split point `k`. Inside this loop, we iterate up to `N` times again to build the prime sets, and factorization takes `O(log V)`. · **Space:** O(N * log(V)) per query, where `V` is the maximum value of a number in `nums`. This space is used to store the sets of prime factors for the prefix and suffix.
**Pros:** Simple to understand and implement.; Correctly solves the problem for very small inputs.
**Cons:** Extremely inefficient due to nested loops.; Repeatedly calculates prime factors for the same numbers across different split points and queries.; Will result in a 'Time Limit Exceeded' error for the given constraints.
### Explanation
The core of this method is a nested loop structure for each query. The outer loop iterates through all valid split points `k` from `1` to `n-1`. The inner loops are responsible for calculating the distinct prime counts. To do this, we use a `HashSet` to keep track of unique prime factors. For a given prefix `nums[0...k-1]`, we iterate from `i=0` to `k-1`, find all prime factors of `nums[i]`, and add them to a `prefix_primes` set. The size of this set gives the count. We do the same for the suffix `nums[k...n-1]`. This entire process is repeated for every single query.

```java
// PrimeHandler class for Sieve and factorization is assumed to exist.
PrimeHandler primeHandler = new PrimeHandler();

// Inside the main method, for each query:
// nums is updated with queries[i][1] at index queries[i][0]

int maxScore = 0;
int n = nums.length;
for (int k = 1; k < n; k++) {
    // Calculate distinct primes in prefix nums[0...k-1]
    Set<Integer> prefixPrimes = new HashSet<>();
    for (int i = 0; i < k; i++) {
        prefixPrimes.addAll(primeHandler.getPrimeFactors(nums[i]));
    }

    // Calculate distinct primes in suffix nums[k...n-1]
    Set<Integer> suffixPrimes = new HashSet<>();
    for (int i = k; i < n; i++) {
        suffixPrimes.addAll(primeHandler.getPrimeFactors(nums[i]));
    }

    maxScore = Math.max(maxScore, prefixPrimes.size() + suffixPrimes.size());
}
// maxScore is the answer for the current query.
```
### Algorithm
- Pre-compute prime factors for all numbers up to the maximum possible value using a Sieve of Eratosthenes.
- For each query `[idx, val]`:
  1. Update the array: `nums[idx] = val`.
  2. Initialize a variable `max_score` to 0.
  3. Iterate through all possible split points `k` from `1` to `n-1`.
  4. For each `k`:
     a. Calculate the number of distinct prime factors in the prefix `nums[0...k-1]`. This involves creating a set, iterating through the prefix elements, finding their prime factors, and adding them to the set. The size of the set is the count for the prefix.
     b. Similarly, calculate the number of distinct prime factors for the suffix `nums[k...n-1]`.
     c. The score for the split `k` is the sum of the two counts.
     d. Update `max_score = max(max_score, score)`.
  5. After checking all `k`, the `max_score` is the result for the current query. Add it to the results array.

## Optimized Calculation per Query
This approach improves upon the brute-force method by optimizing the calculation of scores for all split points. Instead of re-calculating from scratch for each split `k`, it pre-computes the distinct prime counts for all possible prefixes and suffixes in two separate linear passes. After these pre-computations, it finds the maximum score by iterating through the split points one more time.
**Time:** O(Q * N * log(V)). For each query, we perform two passes over the `N` elements to build the count arrays, with factorization at each step. This is still too slow for the given constraints. · **Space:** O(N * log(V)), primarily for storing the prime factors of all numbers and the prefix/suffix count arrays.
**Pros:** Significantly faster than the pure brute-force approach.; Avoids redundant computations within a single query processing.
**Cons:** Still too slow for the given constraints as it performs linear scans for each query.; The prefix and suffix count arrays are completely recomputed for every query, which is inefficient.
### Explanation
For each query, after the array update, we first build an array `prefix_counts`. `prefix_counts[k]` will store the number of distinct prime factors in `nums[0...k-1]`. This can be computed efficiently in a single pass from left to right, maintaining a running `HashSet` of primes. Similarly, we build a `suffix_counts` array where `suffix_counts[k]` stores the count for `nums[k...n-1]`, which is done with a pass from right to left. Once we have these two arrays, we can find the score for any split `k` in `O(1)` time by summing `prefix_counts[k]` and `suffix_counts[k]`. A final loop over `k` finds the maximum possible score.

```java
// PrimeHandler class is assumed to exist.
PrimeHandler primeHandler = new PrimeHandler();

// Inside the main method, for each query:
// nums is updated.
int n = nums.length;

// Pre-calculate all prefix counts
int[] prefixCounts = new int[n + 1];
Set<Integer> currentPrefixPrimes = new HashSet<>();
for (int k = 1; k <= n; k++) {
    currentPrefixPrimes.addAll(primeHandler.getPrimeFactors(nums[k - 1]));
    prefixCounts[k] = currentPrefixPrimes.size();
}

// Pre-calculate all suffix counts
int[] suffixCounts = new int[n + 1];
Set<Integer> currentSuffixPrimes = new HashSet<>();
for (int k = n - 1; k >= 0; k--) {
    currentSuffixPrimes.addAll(primeHandler.getPrimeFactors(nums[k]));
    suffixCounts[k] = currentSuffixPrimes.size();
}

// Find the best split
int maxScore = 0;
for (int k = 1; k < n; k++) {
    // Score for split at k is count for prefix nums[0...k-1] and suffix nums[k...n-1]
    int score = prefixCounts[k] + suffixCounts[k];
    maxScore = Math.max(maxScore, score);
}
// maxScore is the answer for the current query.
```
### Algorithm
- Pre-compute prime factors using a Sieve.
- For each query `[idx, val]`:
  1. Update `nums[idx] = val`.
  2. Create a `prefix_counts` array of size `n+1`. Populate it by iterating from `k=1` to `n`. For each `k`, calculate the distinct prime count in `nums[0...k-1]` by adding the prime factors of `nums[k-1]` to a running set of primes.
  3. Create a `suffix_counts` array of size `n+1`. Populate it by iterating from `k=n-1` down to `0`. For each `k`, calculate the distinct prime count in `nums[k...n-1]` by adding factors of `nums[k]` to a running set.
  4. Initialize `max_score = 0`.
  5. Iterate `k` from `1` to `n-1`:
     a. The score for the split is `prefix_counts[k] + suffix_counts[k]`.
     b. Update `max_score = max(max_score, score)`.
  6. Store `max_score` as the answer for the query.

## Segment Tree with Lazy Propagation
This highly efficient approach reframes the problem to leverage a powerful data structure. The key insight is that the score for a split `k`, `count(prefix) + count(suffix)`, is equal to `count(total) + count(intersection)`. Since `count(total)` (the total number of distinct primes in the array) is independent of the split point `k`, we only need to maximize the size of the intersection, i.e., the number of primes that appear in both the prefix and the suffix. A prime `p` is in the intersection if its first occurrence is before `k` and its last is at or after `k`. This property allows us to use a segment tree to maintain the intersection counts for all `k`. An update to the array only affects a small number of primes, leading to a few efficient updates on the segment tree.
**Time:** O(V log log V + (N + π(V) + Q * log V) * log N). Sieve is `O(V log log V)`. Initialization involves `O(N log V)` for factorization and `O(π(V) * log N)` for populating the segment tree (where `π(V)` is the number of primes up to `V`). Each of the `Q` queries takes `O(log V * log N)` because we update for a few primes (`~log V`), and each update involves `TreeSet` operations (`log N`) and segment tree updates (`log N`). · **Space:** O(V + N log(V)). `O(V)` for the Sieve. The `occurrences` map stores each appearance of a prime factor, totaling `O(N log V)` space in the worst case. The segment tree requires `O(N)` space.
**Pros:** Very efficient, fast enough to pass the given constraints.; Updates are handled logarithmically, not linearly.
**Cons:** Significantly more complex to understand and implement correctly.; Requires careful handling of edge cases and state changes for the `occurrences` map and segment tree.
### Explanation
We maintain the first and last occurrence indices for each prime. A prime `p` with first occurrence `f` and last `l` contributes to the intersection (overlap) for all splits `k` where `f < k <= l`. This means for each such prime, we can think of it as adding `+1` to an 'overlap score' for all `k` in the range `[f+1, l]`. A segment tree with lazy propagation is perfect for this: it can handle range additions and find the maximum value in a range efficiently.

When `nums[idx]` is updated, only the prime factors of the old and new values are affected. For each such prime, we find its `first` and `last` occurrences before the change, undo its contribution on the segment tree (a range update of `-1`), update its set of occurrences, and then apply its new contribution (a range update of `+1` based on new `first` and `last` indices). The total number of distinct primes is also tracked separately. The final answer for a query is this total count plus the maximum overlap score queried from the segment tree.

```java
// Segment Tree with lazy propagation for range add and range max query
class SegmentTree { /* ... standard implementation ... */ }

// Main logic for a query
void processQuery(int idx, int val) {
    int oldVal = nums[idx];
    Set<Integer> oldPrimes = primeHandler.getPrimeFactors(oldVal);
    Set<Integer> newPrimes = primeHandler.getPrimeFactors(val);

    Set<Integer> affectedPrimes = new HashSet<>(oldPrimes);
    affectedPrimes.addAll(newPrimes);

    for (int p : affectedPrimes) {
        TreeSet<Integer> occs = occurrences.get(p);
        // 1. Undo old contribution
        if (occs != null && occs.size() > 1) {
            int oldFirst = occs.first();
            int oldLast = occs.last();
            segTree.rangeUpdate(oldFirst + 1, oldLast, -1);
        }

        // 2. Update occurrences and totalDistinctPrimes
        if (oldPrimes.contains(p)) {
            if (occs.size() == 1) totalDistinctPrimes--;
            occs.remove(idx);
        }
        if (newPrimes.contains(p)) {
            if (occs == null) {
                occs = new TreeSet<>();
                occurrences.put(p, occs);
            }
            if (occs.isEmpty()) totalDistinctPrimes++;
            occs.add(idx);
        }

        // 3. Add new contribution
        if (occs != null && occs.size() > 1) {
            int newFirst = occs.first();
            int newLast = occs.last();
            segTree.rangeUpdate(newFirst + 1, newLast, 1);
        }
    }
    nums[idx] = val;

    int maxOverlap = segTree.queryMax(1, n - 1);
    int result = totalDistinctPrimes + Math.max(0, maxOverlap);
    // Add result to answer list
}
```
### Algorithm
- **Preprocessing:**
  1. Use a Sieve to find the Smallest Prime Factor (SPF) for all numbers up to `10^5`.
  2. Create a map `occurrences` where `occurrences[p]` is a sorted set of indices `i` where `p` is a prime factor of `nums[i]`.
  3. Calculate `total_distinct_primes`, the total number of unique primes across the entire array.
- **Data Structure:**
  1. Use a Segment Tree with Lazy Propagation over the range of split points `k` from `1` to `n-1`. This tree will support range additions and range maximum queries.
- **Initialization:**
  1. The score for a split `k` is `total_distinct_primes + |P_prefix ∩ P_suffix|`. We need to maximize the size of the intersection, which we call `overlap(k)`.
  2. A prime `p` contributes to `overlap(k)` if its first occurrence `f` is `< k` and its last occurrence `l` is `>= k`. This is true for `k` in the range `[f+1, l]`.
  3. For each prime `p` that appears more than once, perform a range update on the segment tree, adding `1` to the range `[f+1, l]`.
- **Per Query:**
  1. For an update `nums[idx] = val`, identify the prime factors of the old and new values (`P_old`, `P_new`).
  2. For each prime `p` in `P_old ∪ P_new`:
     a. **Undo:** If `p` had multiple occurrences before, find its old `first` and `last` indices (`f`, `l`) and subtract `1` from the range `[f+1, l]` in the segment tree.
     b. **Update:** Modify `occurrences[p]` by removing/adding `idx`. Update `total_distinct_primes` if the set of occurrences for `p` becomes empty or non-empty.
     c. **Redo:** If `p` has multiple occurrences now, find its new `first` and `last` indices (`f'`, `l'`) and add `1` to the range `[f'+1, l']` in the segment tree.
  3. After processing all affected primes, query the segment tree for the maximum value, `max_overlap`.
  4. The answer for the query is `total_distinct_primes + max_overlap`.