ABC127 E - Cell Distance - あっとのTECH LOG

問題原文

問題要旨

$H × W$ のグリッドがある。このグリッドに区別のつかない $K$ 個のコマを、各マスにつき最大１つ配置する。
また $K$ 個のコマの、ある配置の仕方についてのコストを、 $\sum_{i=1}^{K-1}\sum_{j=i+1}^{K}(|h_i - h_j| + |w_i - w_j|)$ と定める。
全ての配置の仕方についてのコストの総和を求めよ。ただし配置が異なるとは、片方の配置ではコマがあるが、もう一方にはないようなものをいう。

$1 \leq H × W \leq 2×10^{5}$
$1 \leq K \leq H × W$

解法

いくつかのステップに分かれてます。

制約が完全に罠 $1 \leq H, W \leq 2×10^{5}$ ではないので気をつけましょう。いや罠すぎる。
$H$ と $W$ を分けて考える
こういうのは $H$ と $W$ について分けて考えるのがおすすめです。最後にそれぞれのコストの総和を足し合わせればよいです。
ここからは $W$ 、つまり横方向についてのみ考えていきます。
とりあえずある行にコマを２個おいてみる
とりあえずおいてみます。

とりあえずおいてみる

黄色マスがコマを配置すると決めたところです。この２つのマス間の距離は３ですが、これは何回数えられるでしょうか？
答えは、 $W$ マスからすでに２マス使っていて、加えて $K$ コマからすでに２つ使っているので、 ${}_{W - 2}C_{K - 2}$ 通りになります。
でも実際には縦に $H$ 行分グリッドが続くので、 ${}_{H×W - 2}C_{K - 2}$ 通りになります。
マス間の距離が $d$ になるようなものは何通りあるか？
「ある２つのマスを決めた時に、そのコストが何回カウントされるか」はわかりましたが、まだ間に合いません。
二つのマス間の距離が同じであるならば、それらはまとめて計算できるので、今度は「マス間の距離が $d$ になるようなものは何通りあるか？」を考えていきます。
これは、ある１行に着目した場合には $(W - d)$ 個になります。（ $d$ 個先にマスが存在するようなものは、 $W - d$ マス目までしかありません）
でもこれも３と同じように、縦に $H$ 行分グリッドが続くのでそれらを考慮してやる必要があります。
実際には、必ずしも同じ行に２コマを配置する必要がないので２行の選び方が $H^{2}$ 通りになって、結局マス間の距離が $d$ になるようなものは、 $(W - d)×H^{2}$ 個あります。これらの計算を各 $d$ について行っても、 $0\ leq d \leq W - 1$ なので、十分間に合います。

これで、 $W$ についてのコストの総和を求めることができました。同じように $H$ についても計算して、足し合わせたものが答えになります。

実装

逆元コンビネーションは準備しておいてください。

H, W, K = map(int, input().split())
MOD = 10 ** 9 + 7

ans = 0
# Hについて
for d in range(1, H):
    ans += d * (H - d) * pow(W, 2, MOD) * nCr(H * W - 2, K - 2) % MOD
    ans %= MOD
# Wについて
for d in range(1, W):
    ans += d * (W - d) * pow(H, 2, MOD) * nCr(H * W - 2, K - 2) % MOD
    ans %= MOD

print(ans % MOD)

感想

HとWを分けて考える というのはとても大事なとっかかりですね。
あとは丁寧に考察を積み上げれば解ける、解いてて楽しい問題だと思います。
ただし制約、テメーはダメだ。