ABC156 E - Roaming - あっとのTECH LOG

問題原文

atcoder.jp

問題要旨

区別可能な $N$ 個の部屋があり、最初各部屋には人が１人いる。
この最初の状態から、「ある人が今いる部屋から別の部屋に移動する」という操作が $K$ 回行われた時、各部屋にいる人数の組み合わせとしてありうるものは何通りあるか？

$3 \leq N \leq 2×10^{5}$
$2 \leq K \leq 10^{9}$

解法

部屋が３つ以上、操作が２回以上という制約から、 $N$ 回操作すれば任意の人を任意の場所に配置できることがわかる。
よって、 $K = min(N, K)$ としてよい。これで $K$ が $O(K)$ で扱える制約になった。

次に、 $K$ の制約を一旦無視して、「区別のできない $N$ 人の人を、区別のできる $N$ 個の部屋に割り当てる場合の通り数」を考える。
これは重複組み合わせで解けるから、 ${}_{N + N - 1}C_N$ となる。

ここから、 $K$ 回以上操作しなければ達成できないパターンを引いていく。
まず、 最終的に最低何回操作が行われたのかは、最終的に0人になっている部屋の数 であることに頑張って気づく。（あってるよね）

とすると、0人になっている部屋が $x (K + 1 \leq x \leq N - 1)$ 個あるパターンがダメといえる。
で、ある $x$ のパターンの通り数は、まず「どの部屋が0人になっているか？」で、 ${}_NC_x$ 通りと、「１人以上いることになる区別可能な部屋に $x$ 人を割り当てる通り数」（これも重複組み合わせ）で、 ${}_{N - x + x - 1}C_x$ 通りを掛け合わせた値になる。

実装

逆元コンビネーションでやる。

N, K = map(int, input().split())
K = min(2 * N, K)
MOD = 10 ** 9 + 7

ans = nCr(N + N - 1, N)
for x in range(K + 1, N):
    ans -= nCr(N, x) * nCr(N - x + x - 1, x) % MOD
    ans %= MOD

print(ans % MOD)

感想

0人になっている部屋の数が $K$ 個以上あれば、作り得ない ことに気づけなかった。。。！
意識の反転、視点の反転のトレーニングを。。。