AGC023 - C問題: Painting Machines - 日々ｄｒｄｒする人のメモ

問題: C - Painting Machines

横一列に並んだN個のマスを、N-1台のマシンで塗っていく。i番目のマシンはi番目とi+1番目のマスを黒く塗る。
ある順列Pを与え、その番号の順番でマシンを動かしていき、N個のマスが初めて黒く濡れた時点のPのindex(1-indexed)がスコアとなる。
全ての順列についてこのスコアを求め、その総和を計算する問題。

解けなかった。解法メモ。

この問題の方針は、各スコアKについて、全ての順列の内スコアKになるような通り数を求め、その合計を計算することで解を求める。

コンテスト中もこの方針で計算したが、 ${\sum_{K} (K-1)!(N-K-1)!×F_N(K)}$ となる ${F_N}$ を計算できず積んだ。
この ${F_N(K)}$ では、 ${{}_{K-1}C_{N-K-1}}$ 通り存在するスコアKの各順列について、最後の取れるマシンの番号の通り数を求め、それらの和を取る関数を求めたかった。(最後に取れないマシンの番号iは、番号i以外のK-1台のマシンでi番目とi+1番目が既に塗られている場合である)

想定解法は、K以下のスコアを取る順列の通り数( ${C(K)}$ )を求め、
$\sum_{K= \lceil \frac{N}{2} \rceil}^{N-1} C(K) - C(K-1)$ (mod ${10^9 + 7}$ )
を解として計算する。

${C(K)}$ は ${(K-1)! (N-K-1)! {}_{K-1}C_{N-K-1} = \frac{K!(K-1)!}{(2K-N)!}}$ で計算できる。
この通り数は、稼働するK台のマシンの間の(K-1)箇所に、稼働しないN-K-1台のマシンが高々1台ずつ入り、稼働しないマシンと稼働しないマシンの並べ方を考慮に入れることで計算できる。(どの番号のマシンが稼働するかは、この時点で考慮されている)

簡単な計算だが、逆元の計算含め ${O(N \log N)}$ で解を求めるとPythonではTLEするため、工夫して逆元の計算回数を減らし ${O(N)}$ に落とす必要がある。

下記のように考えれば ${O(N + \log N)}$ で計算できる。下線部もループ処理で逆元なしで計算できる。
${C(K) = \frac{1}{(N-2)!} K! (K-1)! \underline{\frac{(N-2)!}{(2K-N)!}}}$

提出コード (Python3): Submission #2434722 - AtCoder Grand Contest 023

N = int(input())
MOD = 10**9 + 7
 
fact = [1]*(N+1)
 
for i in range(1, N+1):
    fact[i] = r = i*fact[i-1] % MOD
 
cnts = [0]*(N+1)
rev = 1 # rev: 上記の下線部部分
for K in range(N-1, (N+1)//2-1, -1):
    cnts[K] = fact[K]*fact[K-1]*rev % MOD
    rev = rev * (2*K-N) * (2*K-N-1) % MOD
 
ans = 0
for K in range((N+1)//2, N):
    ans += (cnts[K] - cnts[K-1]) * K % MOD
    ans %= MOD
# 最後に一回逆元計算
ans = (ans * pow(fact[N-2], MOD-2, MOD)) % MOD
print(ans)