けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 102 E - Stop. Otherwise... (橙色, 700 点)

数え上げ問題 包除原理 戻すDP 二項係数 DP 重複組合せ AtCoder700点 AtCoder ARC-E 数列 グラフ・盤面・数列の個数の数え上げ 橙色diff 「総和=K」を扱う

いろんな解法がありそうなんな

問題へのリンク

問題概要 (ARC 102 E)

整数 K, N が与えられる。各 x = 2, 3, 4, \dots, 2K に対して、

  • 1 \le a_1 \le a_2 \le \dots \le a_N \le K
  • どの i, j に対しても a_i + a_j \neq x

を満たす整数組 a_1, a_2, \dots, a_N の個数を 998244353 で割った余りを求めよ ( K 面サイコロを N 個振るという設定)。

制約

  • 1 \le K \le 2000
  • 2 \le N \le 2000

解法 1: 漸化式を立ててそれを展開 (本番で通した解法)

例えば x = 5 の場合は、

  • 1 と 4 は同時には出現しない
  • 2 と 3 は同時には出現しない

という条件とみなすことができる。 x がどの場合でも、

  • 1 と b_1 は同時には出現しない
  • 2 と b_2 は同時には出現しない
  • ...
  • p と b_p は同時には出現しない

という形の条件になる。ここで、 b_i = i となるのは構わないが、 i \neq j \le p に対して b_i = j となってはいけない。 p の値は x に応じて決まる。さて、

f(k, p, n) := k 面サイコロを n 個振るときに、「 i が 1 回出たらその後 b_i は出ない」というタイプの条件が p 個付されている状況下での場合の数

として、 f(k, p, n) を再帰的に表すことを考える (注意: 多分ここから先は解法 2 の包除原理の方がわかりやすそう)。

  • a_1 = 1 のとき、それ以後 x 以外の k-1 個しか使えないので f(k-1, p-1, n-1) 通り
  • a_1 = 2 のとき、それ以後 k-2 個しか使えないので f(k-2, p-2, n-1) 通り
  • ...
  • a_1 = p のとき、 f(k-p, 0, n-1) 通り
  • それ以外のとき、 k-p 個から n 個をとる重複組合せなので {}_{k-p}{\rm H}_{n} 通り

これを足す。再帰計算の末端は

  • f(k, 0, n) = {}_{k}{\rm H}_{n} 通り

という感じ。このままこのメモ化再帰を計算すると、 O(K^{3}) の計算量になってしまうので高速化する必要がある。この漸化式について、 p のところが全部 0 になるまで展開を繰り返すと、各 i について

  • {}_{p}{\rm C}_{i} × {}_{k-p}{\rm H}_{n-i}

を合算したものになっていることがわかる。この計算なら、各 k について O(K) 個の和として計算できるので全体の計算量も O(K^{2}) になる。

参考:

#include <iostream>
using namespace std;

const int MAX = 210000;
const int MOD = 998244353;

long long fac[MAX], finv[MAX], inv[MAX];
void COMinit(){
    fac[0] = fac[1] = 1;
    finv[0] = finv[1] = 1;
    inv[1] = 1;
    for(int i = 2; i < MAX; i++){
        fac[i] = fac[i-1] * i % MOD;
        inv[i] = MOD - inv[MOD%i] * (MOD/i) % MOD;
        finv[i] = finv[i-1] * inv[i] % MOD;
    }
}
long long COM(int n, int k){
    if(n < k) return 0;
    if (n < 0 || k < 0) return 0;
    return fac[n] * (finv[k] * finv[n-k] % MOD) % MOD;
}

int main() {
    COMinit();
    int K, N;
    cin >> K >> N;
    for (int query = 2; query <= K*2; ++query) {
        long long res = 0;
        int Q = query;
        if (query > K + 1) Q = (K+1)*2 - query;
        int a = Q/2;
        for (int i = 0; i <= a; ++i) {
            long long tmp = COM(a, i) * COM(K - a + N - i - 1, N - i) % MOD;
            res = (res + tmp) % MOD;
        }
        cout << res << endl;
    }
}

解法 2: 包除原理

僕は上記の思考過程で f(k, p, n) は i = 0, 1, \dots, p についての

  • {}_{p}{\rm C}_{i} × {}_{k-p}{\rm H}_{n-i}

の総和であることを導いたが、これに近い式は包除原理で考えると自然に導出できる。 p 個の制約のうち、 i 個 (以上) を破ってしまうものは

  • {}_{p}{\rm C}_{i} × {}_{k}{\rm H}_{N - 2i} 通り

となる。これを i の偶奇によって足し引きすればよい。微妙に異なる式だけど、多分、ちゃんと式変形すると同じになるんだと思う。

#include <iostream>
using namespace std;

const int MAX = 210000;
const int MOD = 998244353;

long long fac[MAX], finv[MAX], inv[MAX];
void COMinit(){
    fac[0] = fac[1] = 1;
    finv[0] = finv[1] = 1;
    inv[1] = 1;
    for(int i = 2; i < MAX; i++){
        fac[i] = fac[i-1] * i % MOD;
        inv[i] = MOD - inv[MOD%i] * (MOD/i) % MOD;
        finv[i] = finv[i-1] * inv[i] % MOD;
    }
}
long long COM(int n, int k){
    if(n < k) return 0;
    if (n < 0 || k < 0) return 0;
    return fac[n] * (finv[k] * finv[n-k] % MOD) % MOD;
}

int main() {
    COMinit();
    int K, N;
    cin >> K >> N;
    for (int query = 2; query <= K*2; ++query) {
        long long res = 0;
        int Q = query;
        if (query > K + 1) Q = (K+1)*2 - query;
        int q = Q/2;
        for (int i = 0; i <= q; ++i) {
            long long tmp = COM(q, i) * COM(K+N-i*2-1, N-i*2) % MOD;
            if (i & 1) res = (res - tmp + MOD) % MOD;
            else res = (res + tmp) % MOD;
        }
        cout << res << endl;
    }
}

解法 3: 戻す DP

まだちゃんと理解していない...

参考: