問題解決の宝石箱

競技プログラミング/数学ネタ置き場

AGC #005-D ~K Perm Counting

日本語につき問題文は略。

agc005.contest.atcoder.jp

包除原理の適用

いくつかルールがあり、その全てorどれか1つを満たす*1方法を数え上げる問題には、包除原理が使える。 この問題の場合、ルールとは

  •  { a_{i} = i+K } または  { a_{i} = i-K } を満たすこと

である。各 i について合計 N 個のルールがあり、0〜N個のルール採用方法のパターン数がそれぞれ計算できれば、包除原理が適用できて問題が解ける。

ルールを満たすパターンの数え上げ

a[0]からa[N-1]まで順番にルールの採否を決めて処理しようとすると、選んだ数の履歴が必要に思える。 ところが、この問題は観察により回避できる。突破口への鍵を探すため、具体例で考えよう。

例えば、N=10, K=2 としてみると、各 i について a[i] のルールは以下の通り。 a[i] の値が i-2 または i+2 の行と一致するとき、「ルール採用」と考える。

i 1 2 3 4 5 6 7 8 9 10
i-2 - - 1 2 3 4 5 6 7 8
i+2 3 4 5 6 7 8 9 10 -

上の表で、例えば a[1] = 3 に決めたとする。すると、a[2], a[3], a[4] までは a[1] での決定にかかわらず自由にルール採否 を決められる。次に a[1] の決定が関わってくるのは a[5] となる。

  • a[1] = 3 を選んでいた => a[5] = 7 or 不採用
  • a[1]でルールを不採用 した => a[5] = 3 or a[5] = 7 or 不採用

という具合。次に a[5] が関係するのは a[9] であり、このとき a[1] での決定は忘れて良い。

選択の依存関係を図にすると一目瞭然だ。

f:id:hama_DU:20161006232445p:plain

↑依存する位置を同色で塗り分けた図。a[1] での選択は a[5] の選択に影響を与えている。

K=2 に限定せず、図を一般化すると以下のようになる。

f:id:hama_DU:20161006233238p:plain

2*Kで割った余りが同じグループが依存関係となる。

DP

DPの状態は、(ルール採用個数, 現在見ている場所, 直前に i+K のルールを採用したかどうか) 。 「現在見ている場所」は 1, 2, 3, ... , N という順番ではなく、mod 2*K 順に処理する。 例えば N=14, K=3 なら (1 -> 7 -> 13) -> (2 -> 8 -> 14) -> (3 -> 9) -> (4 -> 10) -> ... となる。

遷移ルールは次の通り。

  • ルール不採用 => dp[現在のルール採用個数][次の場所][0] += dp[現在のルール採用個数][現在の場所][0 or 1]
  • i-K のルール採用 => dp[現在のルール採用個数+1][次の場所][0] += dp[現在のルール採用個数][現在の場所][0]
  • i+K のルール採用 => dp[現在のルール採用個数+1][次の場所][1] += dp[現在のルール採用個数][現在の場所][0 or 1]

ただし、i+K のルール採用時でも、mod 2*K が異なる位置に進むときは「直前に i+K のルールを採用したかどうか」フラグを 0 にリセットするのを忘れないように。

状態数が  { O(N^2) }、遷移数が  { O(1) } なので合わせて  { O(N^2) } でこの問題が解ける。

解答コード

*1:あるいは、全てを満たさない