Nasser's Prime Slices (Editorial)

Solution

General Idea:

First, precompute all primes up to 10^7 using the Sieve of Eratosthenes.

For a valid slicing to exist, the total number of primes in the array must be divisible by $k$. If not, the answer is $0$.

Otherwise, think of the problem as placing cut points between groups of exactly $k$ primes. Between consecutive groups of $k$ primes, there may be some non-prime elements. Each non-prime element in such a gap can either be in the left subarray or the right subarray, giving a free binary choice per gap element — so a gap of size $g$ contributes a multiplicative factor of g + 1 to the answer.

More precisely, sweep the array left to right, maintaining a count $cnt$ of primes seen so far in the current group and a counter $check$ for the number of ways to place the current cut:

• When a non-prime is seen and \mathit{cnt} = k, increment $check$ (one more valid cut position).

• When a prime is seen and \mathit{cnt} = k, multiply $ans$ by $check$, reset \mathit{check} = 1 and \mathit{cnt} = 0, then increment $cnt$.

At the end, if the last group does not contain exactly $k$ primes, output $0$. Otherwise output \mathit{ans} \bmod (10^9 + 7).

The solution runs in \mathcal{O}\!\left(A \log \log A + \sum n\right) time per test file, where A = 10^7 is the sieve bound and \sum n \le 2 \times 10^5 is the total input size across all test cases.

#include <bits/stdc++.h>
#define ll long long
using namespace std;

const int MAX = 1e7 + 1, mod = 1e9 + 7;
vector<bool> is_prime(MAX, true);

void sieve() {
    is_prime[0] = is_prime[1] = false;
    for (int i = 2; i * i < MAX; i++)
        if (is_prime[i])
            for (int j = i * i; j < MAX; j += i)
                is_prime[j] = false;
}

void solve() {
    int n, k, cnt = 0, check = 1;
    cin >> n >> k;
    ll ans = 1;
    for (int i = 0; i < n; i++) {
        int x; cin >> x;
        if (!is_prime[x]) {
            if (cnt == k) check++;
        } else {
            if (cnt == k) {
                ans = ans * check % mod;
                check = 1; cnt = 0;
            }
            cnt++;
        }
    }
    if (cnt != k) { cout << 0; return; }
    cout << ans;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    sieve();
    int t; cin >> t;
    while (t--) { solve(); cout << '\n'; }
    return 0;
}