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

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

AtCoder AGC 005 D - ~K Perm Counting (赤色, 900 点)

包除原理勉強会の一環。なんか最初、0〜N-1 を 2K で割った余りで類別したところで鎖を作って、そこでのサイズ i のマッチングの個数を数えたりしたのだけど、それをマージするのに NTT を使うとか血迷った^^;

https://beta.atcoder.jp/contests/agc005/submissions/2819343

普通にまとめて DP すればよかったん。

AGC 005 D - ~K Perm Counting へのリンク

問題概要

整数  N, K があたえられる。
{ 1, 2, \dots, N) の順列 ( a_{1}, a_{2}, \dots, a_{N}) のうち、

  • すべての  i に対して abs( a_{i} - i) \neq K

を満たすものの個数を 924844033 で割ったあまりを求めよ。

制約

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

考えたこと

 K = 0 なら撹乱順列だなという気持ちになる。
全体の枠組みとしては、

 n!
 - (1 箇所を考えてそこが「=」になっている場合の数)(n-1)!
 + (2 箇所を考えてそこが「=」になっている場合の数)(n-2)!
 - ...

という感じになる。撹乱順列であれば、

  • 「1 箇所を考えてそこが「=」になっている場合の数」は単純にその 1 箇所を選んで  {}_{n}{\rm C}_{1} 通り
  • 「2 箇所を考えてそこが「=」になっている場合の数」は単純にその 2 箇所を選んで  {}_{n}{\rm C}_{2} 通り

という感じでよかった。今回は各  a_i に対して「=」になるやつが

  •  a_i = i-K, i+K

と 2 通りある場合と 1 通りの場合とがあるから難しい。

でもさらに考察を進めると、(i 箇所を考えてそこが「=」になっている場合の数) というのは下の図のようになる。

f:id:drken1215:20180710130554j:plain

連結ではないが「 2K 本のパスで構成されたグラフ」というのは「森」なので、そのグラフ上での「サイズ i のマッチングの個数」というのは DP で求めることができる。最初この部分で、

  • パス 1 本 1 本について DP で求めておいて
  • パス  2K 本についての答えを NTT で畳み込んで求める

という風にしてしまったけど、そんな大げさなことをする必要はなかった...

計算量は DP 部分で  O(N^{2})。全体としても  O(N^{2})。editorial によると  O(N \log{N}) でできるらしいが...

#include <iostream>
#include <vector>
#include <cstring>
using namespace std;

const int MOD = 924844033;
const int MAX = 5100;

long long fac[MAX];
void FACInit(){
    fac[0] = fac[1] = 1;
    for(int i = 2; i < MAX; i++) fac[i] = fac[i-1] * i % MOD;
}

int dp[4100][2100][2]; // dp[サイズ i までのパスについて][j 個のマッチングを得ていて][最後の辺が「選ばれているかどうか]

void add(int &a, int b) {
    a += b;
    if (a >= MOD) a -= MOD;
}

int main() {
    FACInit();
    int N, K;
    cin >> N >> K;
    
    vector<int> pathsize(K*2, 0); // 各パスのサイズ
    for (int k = 0; k < K*2; ++k) {
        for (int i = k; i < N; i += K*2) {
            if (i - K >= 0) ++pathsize[k];
            if (i + K < N) ++pathsize[k];
        }
    }
    for (int i = 0; i < pathsize.size(); ++i) {
        if (pathsize[i] == 0)
            pathsize.erase(pathsize.begin() + i--);
    }
    
    // DP: サイズ i のマッチングの個数を求める
    memset(dp, 0, sizeof(dp));
    dp[0][0][0] = 1;
    int kireme = pathsize[0];
    int kireme_ind = 0;
    for (int i = 0; i <= N*2; ++i) {
        for (int j = 0; j <= min(N, i); ++j) {
            add(dp[i+1][j][0], dp[i][j][0]);
            add(dp[i+1][j+1][1], dp[i][j][0]);
            add(dp[i+1][j][0], dp[i][j][1]);
            if (i == kireme) add(dp[i+1][j+1][1], dp[i][j][1]);
        }
        if (i == kireme) {
            ++kireme_ind;
            if (kireme_ind == (int)pathsize.size()) break;
            kireme += pathsize[kireme_ind];
        }
    }

    // 包除原理する
    long long res = 0;
    for (int num = 0; num <= N; ++num) {
        long long matching_size = (dp[kireme][num][0] + dp[kireme][num][1]) % MOD;
        if (num % 2 == 0) (res += fac[N - num] * matching_size % MOD) %= MOD;
        else (res += MOD - fac[N - num] * matching_size % MOD) %= MOD;
    }
    
    cout << res << endl;
}