[ARC168C] Swap Characters 题解

来一篇代码可读性较好的题解。

题意

给你一个长度为 $n$ 的字符串 $S$，字符集为 ABC。定义一次操作：在 $S$ 中自由选择两个字符并交换。执行 $0\sim k$ 次操作，求操作后不同的 $S$ 个数。

对 $998244353$ 取模。

$n\le 2.5\times {10}^5,k\le 100$。

思路

注意到与答案有关的只是 $S$ 中 A、B、C 的个数。

1. 枚举 A 与 B 交换的次数 $i$。
2. 枚举 A 与 C 交换的次数 $j$。
3. 枚举 B 与 C 交换的次数 $k$。
4. 枚举轮换 ABC 的次数 $t$。
5. 用多重组合数直接计算即可。

注意轮换有两种：ABC -> BCA 与 ABC -> CAB。对于 $t\ne 0$ 的情况要分别计算。

时间复杂度 $O(n+k^4)$。AT 神机可以跑过。

代码

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

const int MAXN = 250000 + 5;
const ll MOD = 998244353;

ll n, m;
ll a, b, c;

ll fac[MAXN], invfac[MAXN], inv[MAXN];
void init(ll n) {
    fac[0] = 1; fac[1] = 1;
    invfac[0] = 1; invfac[1] = 1;
    inv[1] = 1;
    for (ll i = 2; i <= n; i++) {
        fac[i] = fac[i-1] * i % MOD;
        inv[i] = (MOD - MOD / i) * inv[MOD % i] % MOD;
        invfac[i] = invfac[i-1] * inv[i] % MOD;
    }
}
ll C(ll n, ll a, ll b, ll c) { // 多重组合数
    if (a < 0 || b < 0 || c < 0) return 0;
    return fac[n] * invfac[a] % MOD * invfac[b] % MOD * invfac[c] % MOD;
}

int main() { ios::sync_with_stdio(0); cin.tie(0);
    cin >> n >> m; init(250000);
    for (int i = 1; i <= n; i++) {
        char ch; cin >> ch;
        if (ch == 'A') a++;
        else if (ch == 'B') b++;
        else if (ch == 'C') c++;
    }
    ll ans = 0;
    for (ll i = 0; i <= m; i++) // AB 交换的次数
        for (ll j = 0; i + j <= m; j++) // AC 交换的次数
            for (ll k = 0; i + j + k <= m; k++) // BC 交换的次数
                for (ll t = 0; i + j + k + t * 2 <= m; t++) { // ABC 轮换的次数
                    ll prod1
                    = C(a, a-i-j-t, i, j+t)
                    * C(b, i+t, b-i-k-t, k) % MOD
                    * C(c, j, k+t, c-j-k-t) % MOD;
                    ll prod2
                    = C(a, a-i-j-t, i+t, j)
                    * C(b, i, b-i-k-t, k+t) % MOD
                    * C(c, j+t, k, c-j-k-t) % MOD;
                    if (t)
                        (ans += prod1 + prod2) %= MOD;
                    else
                        (ans += prod1) %= MOD;
                }
    cout << ans << '\n';
    return 0;
}