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 箇所を考えてそこが「=」になっている場合の数) というのは下の図のようになる。

連結ではないが「 $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;
}