[AtCoder] ABC156 E – Roaming (500点)

2020年2月24日AtCoder剰余, 二項係数, 操作, 数え上げ, 500点, 重複組合せ

問題概要

n 個の部屋に 1 ずつ人がいる。以下を k 回行ったとき、考えられる状態は何通りあるか \(10^9+7\) で割った余りで答えよ。

  • ある部屋 \(i\) にいた人が、\(i \neq j\) を満たす任意の部屋 \(j\) に移動する

問題へのリンク

制約

  • \(3 \leq n \leq 2 \times 10^5\)
  • \(2 \leq k\leq 10^9\)

考え方

何から考えたら良いか分からなくなりますが、「操作を k 回行う」ような問題では、操作で変化する量や変化しない量に注目すると見通しが立つことがあります。

今回では

  • 移動をしても全体の人数は変わらない
  • 移動をすると生じる可能性のある 0 人の部屋の数が増える

などが言えます。

前者は当たり前ですが、後者はどういうことかを説明しましょう。

例えば、はじめ3部屋あったとしたら、

  • 0回移動: (1,1,1) で0人の部屋は最大でも0個
  • 1回移動: (0,1,2) (0,2,1) (1,0,2) (1,2,0) (2,0,1) (2,1,0) で0人の部屋は最大でも1個
  • 2回移動: (1,1,1) (0,1,2) (0,2,1) (1,0,2) (1,2,0) (2,0,1) (2,1,0) (0,0,3) (3,0,0) で0人の部屋は最大でも2個

というようになります。これの太字部分を見ると分かりますが、「i 回移動した時、i-1 回移動した時に比べて、空き部屋が i 個の場合が増える」ことに気がつけばほぼ正解です。
( (1,1,1) だけは例外で、1回移動の時に含まれていません。そこに気がつくと、制約の \(2\leq k\) が親切に見えてきます。)

また、空き部屋は最大でも n-1 部屋までしか作れないので、k が n 以上になっているときはどれだけ移動しても k=n-1の時の値と変わらないです。

よって、以下の計算をすれば良いです。

  • \(\sum_{i=0}^{\min(n-1,k)}\) 空き部屋が \(i\) 個 ある場合の組み合わせ

解法

  • \(\sum_{i=0}^{\min(n-1,k)}\) 空き部屋が \(i\) 個 ある場合の組み合わせ

が答えになります。空き部屋が \(i\) 個 ある場合が何通りあるかは以下のように計算ができます。

$$ _nC_i \times {}_{n-i}H_i$$

意味としては、

  • \(_nC_i\) : n 個の中から 0人の部屋を i 個選ぶ
  • \( {}_{n-i}H_i\) : 残りの \(n-i\) 個の中から重複ありで \(i\) 人の移動先を決める

ということになります。気持ちとしては、剰余を考えないなら以下のような計算をすれば良いです。

重複組合せ \( {}_{n-i}H_i\) は二項係数を用いて以下のように計算することができます。

$$ {}_{x}H_y = {}_{x+y-1}C_y$$

以上より、二項係数を適切な前処理によって、\(O(1)\) で求めることができるならば、十分高速に計算可能です。

C++での実装例

以下のように、剰余の計算をするときは modint 構造体などを利用すると整数のようにプログラムができます。