Operate on Arrays#12384

Operate on Arrays (Editorial)

$Key Observation$

Each operation removes exactly one prime factor from some $a_{i}$ a_i.

a_{i} = p_{1}^{e_{1}} \cdot p_{2}^{e_{2}} \dots p_{m}^{e_{m}},

a_i = p_1^{e_1} \cdot p_2^{e_2} \cdots p_m^{e_m},

Ifthen applying $c$ operations corresponds to removing $c$ prime factors.

c_{1} + c_{2} + \dots + c_{m} = c, 0 \leq c_{j} \leq e_{j} .

c_1 + c_2 + \cdots + c_m = c, \quad 0 \le c_j \le e_j.

Thus, we distribute $c$ removals among exponents:

The number of resulting values depends only on how many times each prime is removed.

$Precomputation$

Let $p r e [x] [c]$ pre[x][c] be the number of distinct values obtainable from $x$ using exactly $c$ operations.

$Smallest Prime Factor (SPF)$

Compute the smallest prime factor for all numbers up to $1 0^{6}$ 10^6 using a sieve.

$DP Construction$

x = p^{e} \cdot y,

x = p^e \cdot y,

Letwhere $p$ is the smallest prime factor and $y$ is not divisible by $p$ .

We remove $k$ copies of $p$ , where $0 \leq k \leq min (c, e)$ 0 \le k \le \min(c, e), and apply the remaining $c - k$ c-k operations to $y$ .

p r e [x] [c] = k = 0 \sum m i n (c, e) p r e [y] [c - k]

pre[x][c] = \sum_{k=0}^{\min(c,e)} pre[y][c-k]

p r e [1] [0] = 1

pre[1][0] = 1

Base case:

$Bounding Operations$

Since $2^{20} > 1 0^{6}$ 2^{20} > 10^6, any number has at most $19$ prime factors. Hence we cap $c \leq 20$ c \le 20.

$Combining Elements (Knapsack DP)$

Let $d p [i] [j]$ dp[i][j] be the number of distinct arrays formed from the first $i$ elements using exactly $j$ operations.

d p [0] [0] = 1

dp[0][0] = 1

Base case:

d p [i] [j] = c = 0 \sum m i n (j, 20) d p [i - 1] [j - c] \cdot p r e [a_{i}] [c]

dp[i][j] = \sum_{c=0}^{\min(j,20)} dp[i-1][j-c] \cdot pre[a_i][c]

Transition:

j = 0 \sum k d p [n] [j]

\sum_{j=0}^{k} dp[n][j]

Final answer:

$Complexity$

Precomputation: $O (M lo g lo g M)$ O(M \log \log M) for sieve and $O (M \cdot 20)$ O(M \cdot 20) for DP.

Per test case: $O (n \cdot k \cdot 20)$ O(n \cdot k \cdot 20).

Space: $O (M \cdot 20)$ O(M \cdot 20) for precomputation and $O (k)$ for DP.</spoiler>

#include<bits/stdc++.h>

using namespace std;

const int MAX_VAL = 1000000;
const int MAX_OP = 20; 
const int MOD = 998244353; 

int spf[MAX_VAL + 1];
int pre[MAX_VAL + 1][MAX_OP + 1];


void globalPrecomputation() {
    for (int i = 2; i <= MAX_VAL; i++) spf[i] = i;
    for (int i = 2; i * i <= MAX_VAL; i++) {
        if (spf[i] == i) {
            for (int j = i * i; j <= MAX_VAL; j += i) {
                if (spf[j] == j) spf[j] = i;
            }
        }
    }

    pre[1][0] = 1;
    for (int op = 1; op <= MAX_OP; op++) pre[1][op] = 0;
    

    for (int i = 2; i <= MAX_VAL; i++) {
        int p = spf[i];
        int temp = i;
        int e = 0;
        
        while (temp % p == 0) {
            e++;
            temp /= p;
        }

// Note: Avg. value of e for all the values (1...MX) is around ~4, So the loop 2 doesn't really add much of complexity.

        int y = temp; 
        
        for (int op = 0; op <= MAX_OP; op++) { // 1
            pre[i][op] = 0;
            for (int k = 0; k <= e; k++) { // 2
                if (op - k >= 0) {
                    pre[i][op] += pre[y][op - k];
                    if(pre[i][op] >= MOD) pre[i][op] >= MOD;
                }
            }
        }
    }
}

void solve() {
    int n, k;
    cin >> n >> k;
    
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    
    vector<long long> dp(k + 1, 0);
    dp[0] = 1;
    
    for (int i = 0; i < n; i++) {
        vector<long long> next_dp(k + 1, 0);
        
        
        for (int j = 0; j <= k; j++) {
            
            for (int c = 0; c <= MAX_OP; c++) {
                if (j + c > k) break; 
                if (pre[a[i]][c] == 0) break; 
                
                next_dp[j + c] = (next_dp[j + c] + dp[j] * pre[a[i]][c]) % MOD;
            }
        }
        dp.swap(next_dp); 
    }
    
    long long ans = 0;
    for (int j = 0; j <= k; j++) {
        ans = (ans + dp[j]) % MOD;
    }
    
    cout << ans << endl;
}

int main() {

    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    globalPrecomputation();
    
    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    
    return 0;
}

C++