LeetCode 3569. Maximize Count of Distinct Primes After Split

June 2, 2025

weekly contest 452。
想通了其實沒很難，只是很煩，而且 python 又卡常數。

題目

https://leetcode.com/problems/maximize-count-of-distinct-primes-after-split/description/

解法

每次查詢步驟：

把 nums[i] 改成 val
選一個索引 k，把 nums 劃分成 nums[0..k-1], nums[k..N-1] 兩半
使得兩半不同的質數個數相加最大

從特殊到一般，先考慮不帶修改的版本要怎麼切才能最大化。
分類討論某質數 p 對於答案的影響：

若 p 在 nums 中只出現一次，那麼怎麼切都無所謂，因為只會對左右其中一半的個數加 1。
若出現兩次以上，只有切在中間才能使兩邊貢獻都加 1；否則只能讓其中一邊的個數加 1。

問題轉換成：求所有質數 p 對於所有分割點 i 的貢獻。

需有序維護 p 的出現位置，使用 sorted list。
記 contribution[i] 為選擇分割點 i 可獲得的不同質數個數和。

若 p 只出現一次，則所有分割點 i 都可以獲得 1 個數的貢獻。
若 p 出現至少兩次，第一次出現在 l，最後一次出現在 r：

分割點 k = [0..l]，所有 p 都在右半，答案加 1
分割點 k = [l+1..r]，兩半都有 p，答案加 2
分割點 k = [r+1..N-1]，所有 p 都在左半，答案加 1

整理後其實增量看起來就像是 [..,1,1,2,2,1,1..]。
可以先把整段都加 1，然後再把中間段在加 1。

d = defaultdict(SL)
contribution = [0]*N

def calc(p):
    sl = d[p]
    if len(sl) == 1:
        # 只出現一次，每個位置貢獻加 1
        for i in range(N):
            contribution[i] += 1
    elif len(sl) >= 2:
        # 出現至少兩次，最左最右在 l, r
        # 中間段 [l+1..r] 貢獻加 2
        # 其餘兩邊 [0..l], [r+1..N1-] 貢獻加 1
        l, r = sl[0], sl[-1]
        for i in range(N):
            contribution[i] += 1
        for i in range(l+1, r+1):
            contribution[i] += 1

處理完所有不同質數 p 後，答案即 max(contribution)。
時間複雜度 O(N)。

再來討論帶修改的版本。

若按照出現次數的變化分類討論會很麻煩，例如：

0 次變 1 次，所有位置都加 1
1 次變 2 次，只有中間位置會加 1
3 次以上，視第一個 / 最後一個位置是否改變，決定那些位置加 1

光想就很囉嗦。
更簡單的方法是把拆成三個步驟：
先撤銷 p 原本的貢獻，增減 p 的位置，然後重新套用 p 的貢獻。

把剛才的貢獻函數稍微調整，加上權重 val，其中 1 代表套用，-1 代表撤銷。

def calc(p, val):
    sl = d[p]
    if len(sl) == 1:
        # 只出現一次，每個位置貢獻加 1
        for i in range(N):
            contribution[i] += val
    elif len(sl) >= 2:
        # 出現至少兩次，最左最右在 l, r
        # 中間段 [l+1..r] 貢獻加 2
        # 其餘兩邊 [0..l], [r+1..N1-] 貢獻加 1
        l, r = sl[0], sl[-1]
        for i in range(N):
            contribution[i] += val
        for i in range(l+1, r+1):
            contribution[i] += val

每次修改只需分開判斷新舊元素是否為質數，就可以複用貢獻函數。

# 舊元素是質數
if old_val in prime:
    calc(old_val, -1)  # 先撤銷原本的貢獻
    d[old_val].remove(i)
    calc(old_val, 1)  # 再套用新的貢獻
# 新元素是質數
if new_val in prime:  # 先撤銷原本的貢獻
    calc(new_val, -1)
    d[new_val].add(i)
    calc(new_val, 1)  # 再套用新的貢獻

nums[i] = new_val  # 記錄修改

因為查詢次數過多，每次複雜度 O(N) 還需要優化。
calc 需要區間修改、區間查詢最大值，可用線段樹優化成 O(N log N)。

但是因為本題卡 python 常數容易超時，所以整個區間 [0..N-1] 修改的操作改用 p_cnt 紀錄次數。

時間複雜度 O((N+Q) log N)。
空間複雜度 O(N)。

MX = 10 ** 5 + 5
sieve = [True] * (MX + 1)
prime = set()
for i in range(2, MX + 1):
    if sieve[i]:
        prime.add(i)
        for j in range(i * i, MX + 1, i):
            sieve[j] = False


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

        d = defaultdict(SL)
        seg = SegmentTree(N)
        p_cnt = 0

        def calc(p, val):
            # 計算質數 p 對各位置的貢獻
            # val = 1/-1 為套用/撤銷
            nonlocal p_cnt
            sl = d[p]
            if len(sl) == 1:
                # 只出現一次，每個位置貢獻加 1
                p_cnt += val
            elif len(sl) >= 2:
                # 出現至少兩次，最左最右在 l, r
                # 中間段 [l+1..r] 貢獻加 2
                # 其餘兩邊 [0..l], [r+1..N-1] 貢獻加 1
                l, r = sl[0], sl[-1]
                p_cnt += val
                seg.update(1, 0, N-1, l+1, r, val)

        # 維護質數位置
        for i, x in enumerate(nums):
            if x in prime:
                d[x].add(i)

        # 計算初始貢獻
        for p in d.keys():
            calc(p, 1)

        ans = []
        for i, new_val in queries:
            old_val = nums[i]
            # 舊元素是質數
            if old_val in prime:
                calc(old_val, -1)  # 先撤銷原本的貢獻
                d[old_val].remove(i)
                calc(old_val, 1)  # 再套用新的貢獻
            # 新元素是質數
            if new_val in prime:  # 先撤銷原本的貢獻
                calc(new_val, -1)
                d[new_val].add(i)
                calc(new_val, 1)  # 再套用新的貢獻

            nums[i] = new_val  # 記錄修改
            ans.append(seg.tree[1] + p_cnt)  # 找最大貢獻

        return ans


class SegmentTree:

    def __init__(self, n):
        self.tree = [0]*(n*4)
        self.lazy = [0]*(n*4)

    def op(self, a, b):
        """
        任意符合結合律的運算
        """
        return a if a > b else b

    def push_down(self, id, L, R, M):
        """
        將區間懶標加到答案中
        下推懶標記給左右子樹
        """
        if self.lazy[id]:
            self.tree[id*2] += self.lazy[id]
            self.lazy[id*2] += self.lazy[id]
            self.tree[id*2+1] += self.lazy[id]
            self.lazy[id*2+1] += self.lazy[id]
            self.lazy[id] = 0

    def push_up(self, id):
        """
        以左右節點更新當前節點值
        """
        self.tree[id] = self.op(self.tree[id*2], self.tree[id*2+1])

    def update(self, id, L, R, i, j, val):
        """
        區間更新
        對[i, j]每個索引都增加val
        """
        if i <= L and R <= j:  # 當前區間目標範圍包含
            self.tree[id] += val
            self.lazy[id] += val
            return
        M = (L+R)//2
        self.push_down(id, L, R, M)
        if i <= M:
            self.update(id*2, L, M, i, j, val)
        if M < j:
            self.update(id*2+1, M+1, R, i, j, val)
        self.push_up(id)

mocowcow

題目

解法