Прості зрізи Нассера (Розбір)

Розв'язок

Загальна ідея: Спочатку попередньо обчисліть усі прості числа до 10^7 за допомогою решета Ератосфена.

Щоб існував допустимий розріз, загальна кількість простих чисел у масиві має ділитися на $k$. Якщо ні, відповідь дорівнює $0$. В іншому випадку розглядайте задачу як розміщення точок розрізу між групами рівно з $k$ простих чисел. Між послідовними групами з $k$ простих чисел можуть бути деякі непрості елементи. Кожен непростий елемент у такому проміжку може належати або лівому підмасиву, або правому, що дає вільний бінарний вибір для кожного елемента проміжку — тому проміжок розміру $g$ вносить мультиплікативний множник g + 1 до відповіді.

Точніше, проходьте масив зліва направо, підтримуючи лічильник $cnt$ простих чисел, побачених у поточній групі, та лічильник $check$ для кількості способів розмістити поточний розріз:

• Коли зустрічається непросте число і \mathit{cnt} = k, збільшіть $check$ на $1$ (ще одна допустима позиція розрізу).

• Коли зустрічається просте число і \mathit{cnt} = k, помножте $ans$ на $check$, скиньте \mathit{check} = 1 і \mathit{cnt} = 0, потім збільшіть $cnt$ на $1$.

Наприкінці, якщо остання група не містить рівно $k$ простих чисел, виведіть $0$. Інакше виведіть \mathit{ans} \bmod (10^9 + 7).

Розв'язок працює за час \mathcal{O}\!\left(A \log \log A + \sum n\right) на тестовий файл, де A = 10^7 — межа решета, а \sum n \le 2 \times 10^5 — загальний розмір вхідних даних по всіх тестових випадках.

#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;
}