题解:AT_agc046_c [AGC046C] Shift

· · 题解

突然懂了计数 dp 不一定是要模拟过程而是可以刻画最终结果。

下文中 m 代表题目中的 k,k 代表分段后的段数,n 代表 |S|。

这道题的操作就是把后面的一个 1 放到一个 0 的前面来。

我们先把序列分成极长 01 段。如果把后面的一个 1 放入这个段里面然后再把这个 1 放到前面去,这个操作等价于直接把后面这个 1 放到前面去,所以我们只用考虑一个数只移动一次的情况即可。所以最多操作 \max(n,m) 次。

这样我们就可以直接 dp 了,令 dp_{i,j,c} 表示第 i 个段,插进来了 j 个 1,送出去了 c 个 1。考虑转移

考虑第 i 段,它原有 siz_i 个 1。分三种情况:

  1. 不变:dp_{i,j,c} \gets dp_{i-1,j,c}

  2. 移入 p 个:dp_{i,j,c} \gets \sum_{p=1}^{j-c} dp_{i-1,j-p,c}

  3. 移出 p 个:dp_{i,j,c} \gets \sum_{p=1}^{\min(c,siz_i)} dp_{i-1,j,c-p}

最终所有段处理完后,全局移入总数必须等于全局移出总数,又由于每个 1 最多移动一次,移动总数不超过 n,故令 m \gets \min(m, n)。答案为:\sum_{i=0}^{m} dp_{k,i,i}

暴力转移是 O(n^4) 的,在原 OJ 上是可以通过的。可以用前缀和优化做到 O(n^3)。

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int mod = 998244353; 
string s;
const int N = 305;
int dp[N][N][N];
int siz[N],k;
int n,m;
int ans = 0;
signed main(){
    ios::sync_with_stdio(0);
    cin.tie(nullptr);
    cin >> s >> m;
    n = s.size();
    s = " " + s;
    s.push_back('0');
    int t = 0;
    for(int i = 1;i <= n + 1;i++){
        if(s[i] == '1'){
           t++;
        }else{
            siz[++k] = t;
            t = 0;
        }
    }
    dp[0][0][0] = 1;
    m = min(m,n);
    for(int i = 1;i <= k;i++){
        for(int j = 0;j <= m;j++){
            for(int c = 0;c <= j;c++){
                dp[i][j][c] = dp[i - 1][j][c];
                for(int p = 1;p <= j - c;p++){
                    dp[i][j][c] = (dp[i][j][c] + dp[i - 1][j - p][c]) % mod; 
                }
                for(int p = 1;p <= min(c,siz[i]);p++){
                    dp[i][j][c] = (dp[i][j][c] + dp[i - 1][j][c - p]) % mod;
                }
            }
        }
    }
    for(int i = 0;i <= m;i++){
        ans = (ans + dp[k][i][i]) % mod;
    }
    cout << ans;
    return 0;
}