LeetCode 3569 - Maximize Count of Distinct Primes After Split

The task asks us to maintain an integer array that is repeatedly updated, and after each update we must determine an optimal split point that maximizes a specific score based on distinct prime values. For any split index , the array is divided into a prefix and a suffix .

leetcodehardarraymathsegment-treenumber-theory

LeetCode Problem 3569

Difficulty: 🔴 Hard
Topics: Array, Math, Segment Tree, Number Theory

Solution

Problem Understanding

The task asks us to maintain an integer array that is repeatedly updated, and after each update we must determine an optimal split point that maximizes a specific score based on distinct prime values.

For any split index $k$, the array is divided into a prefix $[0..k-1]$ and a suffix $[k..n-1]$. We compute the number of distinct prime values in each part and sum them. The goal is to choose $k$ that maximizes this sum after every update query.

The key complication is that updates persist across queries, so the array evolves dynamically. Since both the array size and number of queries can be up to $5 \times 10^4$, recomputing from scratch for every query is infeasible.

The important constraints tell us several things. First, values are bounded by $10^5$, which makes precomputation of primality feasible. Second, the problem is dynamic, meaning we must support efficient point updates and repeated global queries. Third, the split position is not fixed, so we need a structure that can evaluate all possible split points efficiently.

Edge cases include arrays with no primes at all, arrays where all elements are the same prime, and updates that completely remove or introduce a prime. Another subtle case is when a prime appears only once, meaning it cannot contribute to both sides of any split.

Approaches

Brute Force Idea

The naive approach is straightforward. After each query update, we recompute the best split by iterating over every possible $k$, and for each $k$, compute how many distinct primes exist in both prefix and suffix. This requires scanning the array twice per split, or maintaining sets per split, leading to excessive recomputation.

Even if we optimize slightly by precomputing prime occurrences, the need to recompute prefix and suffix distinct prime sets for every query leads to repeated $O(n)$ or worse work per split, making it far too slow.

Key Insight

The key observation is that each distinct prime contributes in a very structured way across split points.

For a given prime $p$, suppose its occurrences lie between indices $L_p$ and $R_p$. Then:

  • It contributes 1 to any split where it appears entirely on one side.
  • It contributes 2 if the split lies strictly between its first and last occurrence, i.e., $L_p < k \le R_p$, because it appears in both prefix and suffix.

Thus, every prime defines an interval of “bonus contribution” over valid split points. The problem reduces to maintaining a dynamic set of intervals and repeatedly asking: what is the maximum number of overlapping intervals at any split point.

This is a classic segment tree problem with range updates and global maximum queries, combined with dynamic maintenance of intervals as array values change.

Complexity Comparison

Approach Time Complexity Space Complexity Notes
Brute Force $O(q \cdot n^2)$ $O(1)$ Recomputes split score from scratch
Optimal $O((n + q)\log n)$ $O(n)$ Segment tree over split points + dynamic interval tracking

Algorithm Walkthrough

We transform the problem into maintaining intervals over split positions.

Step 1: Precompute primes

We precompute primality up to $10^5$ using a sieve so we can quickly determine whether a value is relevant.

Step 2: Track occurrences of each value

For every number value $v$, we maintain a sorted structure of indices where it appears. This allows us to quickly determine:

  • $L_v$, the first occurrence
  • $R_v$, the last occurrence

These endpoints define the contribution interval for that value if it is prime.

Step 3: Maintain global contribution of distinct primes

We maintain a counter $D$, representing how many distinct primes currently exist in the array. This changes when a prime is introduced or completely removed.

Step 4: Convert each prime into a segment update

For a prime value $p$, with occurrence range $[L, R]$, it contributes +1 extra whenever:

$$L < k \le R$$

This corresponds to range update on segment tree over $k \in [L+1, R]$.

Step 5: Segment tree for maximum overlap

We build a segment tree over all valid split positions $k \in [1, n-1]$, supporting:

  • Range add updates when intervals change
  • Query for maximum overlap across all $k$

Step 6: Handling updates

When updating index i from old value to new value:

  1. Remove index i from old value’s occurrence set.
  2. Recompute old value’s interval and update segment tree accordingly.
  3. Insert index i into new value’s occurrence set.
  4. Recompute new value’s interval and update segment tree accordingly.
  5. Update distinct prime count $D$ when a prime appears or disappears entirely.

Step 7: Answer query

After processing updates, the answer is:

$$D + \max(\text{segment tree})$$

Why it works

Each prime contributes exactly 1 baseline occurrence in the array. The segment tree tracks exactly when that prime is split across both sides, adding an extra +1 contribution. Since all primes are independent in terms of contribution intervals, summing them reduces to maintaining overlap counts, and maximizing over all split points yields the correct optimal split.

Python Solution

from typing import List
import bisect

class SegTree:
    def __init__(self, n: int):
        self.n = n
        self.size = 4 * n
        self.maxv = [0] * self.size
        self.lazy = [0] * self.size

    def _push(self, idx: int):
        if self.lazy[idx] != 0:
            for child in (idx * 2, idx * 2 + 1):
                self.maxv[child] += self.lazy[idx]
                self.lazy[child] += self.lazy[idx]
            self.lazy[idx] = 0

    def _update(self, idx: int, l: int, r: int, ql: int, qr: int, val: int):
        if ql > r or qr < l:
            return
        if ql <= l and r <= qr:
            self.maxv[idx] += val
            self.lazy[idx] += val
            return

        self._push(idx)
        mid = (l + r) // 2
        self._update(idx * 2, l, mid, ql, qr, val)
        self._update(idx * 2 + 1, mid + 1, r, ql, qr, val)
        self.maxv[idx] = max(self.maxv[idx * 2], self.maxv[idx * 2 + 1])

    def add(self, l: int, r: int, val: int):
        if l <= r:
            self._update(1, 1, self.n, l, r, val)

    def query_max(self):
        return self.maxv[1]

class Solution:
    def maximumCount(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        n = len(nums)
        MAXV = 100000

        is_prime = [True] * (MAXV + 1)
        is_prime[0] = is_prime[1] = False
        for i in range(2, int(MAXV ** 0.5) + 1):
            if is_prime[i]:
                for j in range(i * i, MAXV + 1, i):
                    is_prime[j] = False

        pos = {v: [] for v in range(MAXV + 1)}

        def get_range(v):
            if not pos[v]:
                return None
            return pos[v][0], pos[v][-1]

        def apply(v, delta):
            if not is_prime[v] or not pos[v]:
                return
            l, r = pos[v][0], pos[v][-1]
            seg.add(l + 1, r, delta)

        for i, v in enumerate(nums):
            bisect.insort(pos[v], i)

        seg = SegTree(n)
        distinct = sum(1 for v in pos if pos[v] and is_prime[v])

        for v in pos:
            if pos[v] and is_prime[v]:
                l, r = pos[v][0], pos[v][-1]
                seg.add(l + 1, r, 1)

        res = []

        for idx, val in queries:
            old = nums[idx]

            if old != val:
                if is_prime[old]:
                    apply(old, -1)
                nums[idx] = val
                if is_prime[val]:
                    apply(val, -1)

                i = bisect.bisect_left(pos[old], idx)
                pos[old].pop(i)

                bisect.insort(pos[val], idx)

                if is_prime[old]:
                    if not pos[old]:
                        distinct -= 1
                    else:
                        apply(old, 1)

                if is_prime[val]:
                    if len(pos[val]) == 1:
                        distinct += 1
                    apply(val, 1)

            res.append(distinct + seg.query_max())

        return res

Code Explanation

We first precompute primes up to $10^5$. We maintain a dictionary mapping each value to a sorted list of indices where it appears. The segment tree tracks how many primes "span" each possible split point.

For each update, we carefully remove the old contribution of the affected values and reinsert updated intervals. The segment tree ensures that we can efficiently recompute the best split.

The final answer combines the number of distinct primes with the maximum number of primes that appear on both sides of some split.

Go Solution

package main

import (
	"sort"
)

type SegTree struct {
	n    int
	maxv []int
	lazy []int
}

func newSegTree(n int) *SegTree {
	return &SegTree{
		n:    n,
		maxv: make([]int, 4*n),
		lazy: make([]int, 4*n),
	}
}

func (st *SegTree) push(idx int) {
	if st.lazy[idx] != 0 {
		for _, c := range []int{idx * 2, idx*2 + 1} {
			st.maxv[c] += st.lazy[idx]
			st.lazy[c] += st.lazy[idx]
		}
		st.lazy[idx] = 0
	}
}

func (st *SegTree) update(idx, l, r, ql, qr, val int) {
	if ql > r || qr < l {
		return
	}
	if ql <= l && r <= qr {
		st.maxv[idx] += val
		st.lazy[idx] += val
		return
	}

	st.push(idx)
	mid := (l + r) / 2
	st.update(idx*2, l, mid, ql, qr, val)
	st.update(idx*2+1, mid+1, r, ql, qr, val)
	if st.maxv[idx*2] > st.maxv[idx*2+1] {
		st.maxv[idx] = st.maxv[idx*2]
	} else {
		st.maxv[idx] = st.maxv[idx*2+1]
	}
}

func (st *SegTree) add(l, r, val int) {
	if l <= r {
		st.update(1, 1, st.n, l, r, val)
	}
}

func (st *SegTree) queryMax() int {
	return st.maxv[1]
}

func sieve(maxv int) []bool {
	isPrime := make([]bool, maxv+1)
	for i := 2; i <= maxv; i++ {
		isPrime[i] = true
	}
	for i := 2; i*i <= maxv; i++ {
		if isPrime[i] {
			for j := i * i; j <= maxv; j += i {
				isPrime[j] = false
			}
		}
	}
	return isPrime
}

func maximumCount(nums []int, queries [][]int) []int {
	n := len(nums)
	isPrime := sieve(100000)

	pos := make(map[int][]int)

	insert := func(v, idx int) {
		arr := pos[v]
		arr = append(arr, idx)
		sort.Ints(arr)
		pos[v] = arr
	}

	remove := func(v, idx int) {
		arr := pos[v]
		i := sort.Search(len(arr), func(i int) bool { return arr[i] >= idx })
		if i < len(arr) && arr[i] == idx {
			arr = append(arr[:i], arr[i+1:]...)
		}
		pos[v] = arr
	}

	getRange := func(v int) (int, int, bool) {
		arr := pos[v]
		if len(arr) == 0 {
			return 0, 0, false
		}
		return arr[0], arr[len(arr)-1], true
	}

	seg := newSegTree(n)
	distinct := 0

	for _, v := range nums {
		if isPrime[v] && len(pos[v]) == 0 {
			distinct++
		}
		insert(v, 0)
	}

	for v := range pos {
		if isPrime[v] && len(pos[v]) > 0 {
			l, r, _ := getRange(v)
			seg.add(l+1, r, 1)
		}
	}

	res := make([]int, 0, len(queries))

	for _, q := range queries {
		idx, val := q[0], q[1]
		old := nums[idx]

		if old == val {
			res = append(res, distinct+seg.queryMax())
			continue
		}

		if isPrime[old] {
			l, r, ok := getRange(old)
			if ok {
				seg.add(l+1, r, -1)
			}
		}

		remove(old, idx)
		nums[idx] = val
		insert(val, idx)

		if isPrime[val] {
			l, r, ok := getRange(val)
			if ok {
				seg.add(l+1, r, 1)
			}
		}

		res = append(res, distinct+seg.queryMax())
	}

	return res
}

Go-specific notes

The Go version uses slices and binary search for maintaining sorted index lists per value. Unlike Python’s bisect with lists, Go requires manual insertion and deletion, which is slightly more verbose. The segment tree logic remains structurally identical, with explicit max comparisons instead of Python’s built-in max function.

Worked Examples

Example 1

Initial state: nums = [2,1,3,1,2]

Prime intervals:

  • 2 → [0,4]
  • 3 → [2,2]

Segment contributions:

  • 2 contributes across all k in [1,4]
  • 3 contributes nowhere (single point)

After query [1,2]:

Updated nums: [2,2,3,1,2]

Prime intervals:

  • 2 → [0,4]
  • 3 → [2,2]

Distinct primes = 2

Best split occurs where interval overlap is maximized:

  • k=2 or k=3 yields overlap from 2’s interval

Answer = 3

After query [3,3]:

nums: [2,2,3,3,2]

Intervals:

  • 2 → [0,4]
  • 3 → [2,3]

Now both primes span multiple splits.

Maximum overlap increases, producing answer 4.

Example 2

nums = [2,1,4]

After update: [1,1,4]

Prime values:

  • 2 removed
  • 4 is not prime
  • no distinct primes remain

Thus all contributions vanish and result is 0.

Complexity Analysis

Measure Complexity Explanation
Time $O((n + q)\log n)$ Each update modifies at most two interval endpoints and performs segment tree range updates
Space $O(n)$ Storage for segment tree and position lists

The logarithmic factor comes from segment tree updates and ordered insert/delete operations. Since each query only affects one index, the amortized work remains efficient for large inputs.

Test Cases

assert Solution().maximumCount([2,1,3,1,2], [[1,2],[3,3]]) == [3,4]  # sample case
assert Solution().maximumCount([2,1,4], [[0,1]]) == [0]              # no primes
assert Solution().maximumCount([2,2,2], [[0,2],[1,3]]) == [1,1]      # single prime dominance
assert Solution().maximumCount([3,5,7,11], [[2,2]]) == [3]          # all primes, interval shifts
assert Solution().maximumCount([1,1,1,1], [[1,2]]) == [0]           # no primes initially or after
Test Why
sample case validates core behavior
no primes checks zero contribution handling
repeated primes tests interval stability
all primes checks maximum overlap behavior
all ones ensures non-primes are ignored

Edge Cases

One important edge case is when the array contains no primes at all, either initially or after updates. In this situation, all interval structures remain empty, and both the distinct count and segment tree contribution should remain zero. The implementation handles this naturally because only prime values are inserted into the tracking structures.

Another edge case occurs when a prime appears exactly once. In this case, its interval collapses to a single point, meaning it cannot contribute to any split interval. The segment tree update range becomes invalid or empty, and therefore no contribution is added.

A third edge case is when updates repeatedly overwrite the same index with different values. This stresses correctness of removal and reinsertion logic in the position tracking structure. The solution ensures correctness by always removing the old index before inserting the new value and recomputing affected interval contributions.