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

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

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

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

問題概要

整数  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

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

参考: