Shuffle 题解

Shuffle

题目大意

给定一个长度为 \(n\) 的 01 序列 \(a\)，你可以进行至多一次以下操作：

• 选定 \(a\) 的一个连续段，满足连续段内恰好有 \(k\) 个 \(1\)，将该连续段任意排列。

问能产生多少种不同的 01 序列。

思路分析

（这题 \(n\) 完全可以开到 \(10^6\) 或是 \(10^7\)，因为存在 \(O(n)\) 的做法。）

考虑 DP。

设 \(f_i\) 表示只考虑 \(1\sim i\) 中的字符，能产生多少种不同的 01 序列。

那么可以列出 DP 方程：

\[f_i=\begin{cases}f_{i-1}+{m-1\choose k-1}&s_i=0\\f_{i-1}+{m-1\choose k}&s_i=1\end{cases}
\]

其中，\(m\) 是 \(i\) 往左的极长 \(k\) 个 \(1\) 的连续段的长度。

解释一下：

我们在考虑 \(f_{i-1}\) 时，是把 \(s_i\) 恒定为 \(s_i\) 做的，而在考虑 \(f_i\) 时为了避免算重，我们强制钦定在 \(i\) 的位置放 \(s_i\) 的相反的数，也就是说，若 \(s_i=0\)，我们强制这个位置放 \(1\)，那么方案数就是 \({m-1\choose k-1}\)，在前 \(m-1\) 个位置上放剩下的 \(k-1\) 个数。\(s_i=1\) 类似。

注意初值，当 \(m\) 第一次存在时，\(f_i={m\choose k}\)，这是因为 \(f_i\) 前面没有人算它的重。

实现方式有很多种，我使用了二分实现，时间复杂度为 \(O(n\log n)\)。

代码

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <cmath>
#include <set>

using namespace std;
const int N = 1001000, L = 1000000, mod = 998244353;
#define inf 0x3f3f3f3f
#define int long long

int n, k, ans;
int fac[N], inv[N], sum[N];

set <int> S;

char s[N];

int q_pow(int a, int b){
    int res = 1;
    while (b) {
        if (b & 1) res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}

int C(int n, int m){
    if (n < m) return 0;
    return fac[n] * inv[n - m] % mod * inv[m] % mod;
}

int find(int s, int k){ // 找到从 s 往左的第 k + 1 个 1 的位置的右边
    if (sum[s] < k) return inf;
    int l = 1, r = s, ans = 1;
    while (l <= r) {
        int mid = (l + r) >> 1;
        if (sum[s] - sum[mid - 1] <= k) r = mid - 1, ans = mid;
        else l = mid + 1;
    }
    return *(--S.lower_bound(ans)) + 1;
}

signed main(){
    fac[0] = 1;
    for (int i = 1; i <= L; i ++) fac[i] = fac[i - 1] * i % mod;
    inv[L] = q_pow(fac[L], mod - 2);
    for (int i = L; i >= 1; i --) inv[i - 1] = inv[i] * i % mod;
    scanf("%lld %lld", &n, &k);
    scanf("%s", s + 1);
    for (int i = 1; i <= n; i ++) sum[i] = sum[i - 1] + (s[i] == '1');
    S.insert(0);
    for (int i = 1; i <= n; i ++) if (s[i] == '1') S.insert(i);
    int flag = 0;
    for (int i = 1; i <= n; i ++) {
        if (sum[i] == k && !flag) ans = (ans + C(i, k)) % mod, flag = 1;
        else ans = (ans + C(i - find(i, k), k - (s[i] == '0'))) % mod;
    }
    if (!flag) ans = 1;
    cout << ans << '\n';
    return 0;
}