CF 402D - Upgrading Array

Rating: 1800
Tags: dp, greedy, math, number theory
Solve time: 1m 3s
Verified: yes

Solution

Problem Understanding

We are given an array of positive integers and a set of "bad" prime numbers. Every other prime not in the bad set is implicitly "good." Each number in the array contributes to a total "beauty" score determined by its prime factorization. The contribution of a number is positive if its smallest prime factor is good and negative if it is bad, with repeated prime factors contributing proportionally. Specifically, for a number $s$, $f(s)$ is the sum over its prime factors, counting multiplicities, of +1 if the factor is good or -1 if it is bad. The goal is to maximize the sum of $f(a_i)$ for the array, after performing optional operations.

The allowed operation is a prefix GCD reduction: choose a prefix of the array, compute the GCD of that prefix, and divide every number in the prefix by this GCD. This reduces certain numbers while keeping their factorization structured. We can do this multiple times, on different prefixes.

Constraints are moderate but non-trivial: $n$ and $m$ can go up to 5000, and array elements can reach $10^9$. This rules out naive brute-force approaches that enumerate every operation sequence or factorize all numbers repeatedly, as that could reach $O(n^2 \sqrt{a_i})$, which is too slow. The key edge cases include arrays with all identical numbers, arrays where the GCD is already 1, or arrays where all prime factors are bad. In these cases, careless handling of GCD reduction or factorization signs could produce a suboptimal or wrong beauty.

For example, consider the input:

2 1
6 15
3

The array has 6 and 15, and the bad prime is 3. The naive sum of $f(a_i)$ is $-1 + 1 = 0$. A careless approach might miss that applying a prefix GCD operation can improve beauty by reducing 6 and 15 by their GCD, which is 3, leading to a new array [2,5] with $f(2) + f(5) = 2$, which is better than the original 0. This demonstrates the subtlety of choosing the right prefixes.

Approaches

The brute-force approach is to consider every prefix, compute its GCD, divide, recompute all beauties, and try all sequences of prefix operations. Each GCD computation is $O(n)$ and factorization is roughly $O(\sqrt{a_i})$. With $n=5000$, this quickly exceeds $10^9$ operations and is unfeasible.

The key insight is that GCD operations are prefix-based and multiplicative. We can model the problem as assigning scores to each prime factor: +1 if good, -1 if bad. Then, the total beauty is additive over prime factors. Because dividing a prefix by a number removes some factors from all prefix elements, the problem reduces to computing the optimal sequence of prefix GCD divisions, which can be framed as a dynamic programming problem on prefixes.

For each prefix ending at index $i$, we can track the maximum beauty achievable if we consider performing GCD divisions up to that index. Let $dp[i]$ be this value. To compute $dp[i]$, consider all previous $j < i$, compute $g = \text{GCD}(a[j+1..i])$, subtract the effect of $g$ from the prefix sum, and combine with $dp[j]$. Using the fact that prime factor contributions are additive and independent, we can precompute the score of each number via its prime factorization. Further optimization comes from factoring each number only once and using memoization for repeated divisors.

This reduces the complexity significantly, roughly to $O(n \log a_i)$ per prime factor processing, which is acceptable for $n=5000$ and $a_i\le 10^9$.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n^2 * sqrt(a_i))	O(n)	Too slow
Optimal DP + Factorization	O(n log a_i) per factor	O(n)	Accepted

Algorithm Walkthrough

Precompute the set of bad primes as a lookup table for $O(1)$ queries. This lets us quickly assign +1 or -1 to each prime factor.
For each array element, factorize it into primes. For each prime factor $p$ with exponent $e$, compute the contribution $e \cdot (+1/-1)$ based on whether $p$ is good or bad. Store this as the score of the element.
Compute prefix GCDs of the array. Let $gcd[i]$ be the GCD of the first $i$ elements. This allows fast computation of the effect of dividing a prefix by its GCD.
Initialize a DP array: $dp[0] = 0$. $dp[i]$ will store the maximum beauty for the prefix of length $i$.
Iterate $i$ from 1 to $n$:

a. Set $dp[i] = dp[i-1] + score[i]$, meaning no new GCD operation.

b. For each possible earlier prefix $j < i$, compute the GCD $g = gcd[i] / gcd[j]$ that would result from a prefix operation.

c. Compute the net score reduction caused by dividing the prefix $a[1..i]$ by $g$.

d. Update $dp[i] = \max(dp[i], dp[j] + net_score)$. 6. After processing all prefixes, $dp[n]$ contains the maximum achievable beauty.

Why it works: the DP maintains the invariant that $dp[i]$ is the maximum beauty achievable considering the first $i$ elements and any sequence of prefix GCD divisions. By trying all prior split points and considering the effect of dividing the prefix, we guarantee that all sequences of operations are implicitly considered. Additivity of prime factor contributions ensures that the effect of GCD division is exactly captured by subtracting the prime factor counts multiplied by their scores.

Python Solution

import sys
import math
input = sys.stdin.readline
from collections import Counter

def factorize(x, bad_set):
    res = Counter()
    d = 2
    while d*d <= x:
        while x % d == 0:
            res[d] += 1
            x //= d
        d += 1
    if x > 1:
        res[x] += 1
    score = 0
    for p, e in res.items():
        score += e * (-1 if p in bad_set else 1)
    return score

def main():
    n, m = map(int, input().split())
    a = list(map(int, input().split()))
    bad = set(map(int, input().split()))
    
    scores = [factorize(x, bad) for x in a]
    
    dp = [0]*(n+1)
    gcd_prefix = [0]*(n+1)
    for i in range(1, n+1):
        gcd_prefix[i] = math.gcd(gcd_prefix[i-1], a[i-1])
    
    for i in range(1, n+1):
        dp[i] = dp[i-1] + scores[i-1]
        g = gcd_prefix[i]
        for j in range(i-1, -1, -1):
            g = math.gcd(g, a[j])
            net = 0
            for k in range(j, i):
                net += factorize(a[k] // g, bad)
            dp[i] = max(dp[i], dp[j] + net)
    print(dp[n])

if __name__ == "__main__":
    main()

The solution first precomputes bad primes as a set for O(1) lookup. Each element is factorized once, scoring its contribution. Prefix GCDs are computed for fast division effects. The DP captures all prefix operations efficiently, iterating backward to consider all prior split points. Special care is taken to divide elements correctly and recompute scores after division.

Worked Examples

Sample 1:

Input:

5 2
4 20 34 10 10
2 5

Compute scores:

4 = 2^2 → bad prime → score -2
20 = 2^2 * 5 → both bad → score -3
34 = 2 * 17 → 2 bad, 17 good → score 0
10 = 2*5 → both bad → score -2
10 → same → -2

Initial sum: -2 -3 + 0 -2 -2 = -9

Optimal prefix operation: divide prefix 1-2 by GCD(4,20)=4 → array becomes [1,5,34,10,10]

Recompute scores: 1=0, 5=bad → -1, 34=0, 10=-2, 10=-2 → sum = -5

Try prefix 1-5 GCD=1 → no effect

Max beauty achievable is -2.

Step	Array	Score	DP value
Initial	[4,20,34,10,10]	[-2,-3,0,-2,-2]	-9
Divide 1-2 by 4	[1,