ABC127E

問題

atcoder.jp $N \times M$ マスに配置する $K$ 個の駒の間のマンハッタン距離の総和を，とりうる配置のパターンでさらに総和する問題です。

解説

公式解説を読んでも理解に苦しめられたので，少しだけ噛み砕いて記事にします。
そしてなんと，今回からMarkdown記法で書くことにしました。ほとんど更新してなくてすみません。

2点 $P(x_i, y_i), Q(x_j, y_j)$ 間のマンハッタン距離を場合の数だけ掛け合わせると良いです。これについて，具体例をもとに３つのSTEPに分けて理解していきたいと思います。

具体例 $(N, M, K) = (2, 2, 3)$ で考える

$2 \times 2$ のマスに3個の点を配置する場合を考えます。

STEP 1

点の配置の仕方は図のようになります。点を区別するために番号を振っています。
各配置におけるマンハッタン距離の総和を４通りの場合で総和をとると答えになります。

STEP 2

すべての配置を列挙するのは難しいので，２点間の距離のみに着目して，それが答えに何回足されるかを考えます。残りの点の配置の仕方は，残りの2マスの中から1マスを選ぶので ${}_{2} C_{1}$ 通りです。一般化すると， ${}_{NM-2} C_{K-2}$ 通りになります。 $N \times M \leqq 2 \times 10^{5}$ の制約では， $O(NM)$ で算出しても間に合います。
PQの選び方を列挙したら，それぞれのマンハッタン距離を算出し，場合の数をかけて合算します。

ここで， $N \times M$ のグリッドから2点を選ぶ場合の数は ${}_{NM} C_{2}$ であり，これを全探索してしまうと計算量は $O(N^{2}M^{2})$ になります。もし，pythonで書くならば，

for i in range(N*M-1):
    for j in range(i+1, N*M):
        ...

のようになると思います。これでは間に合いません。そこで次の考え方に移ります。

STEP 3

マンハッタン距離の計算において， $x$ , $y$ は独立しているので，今回は分離して考えます。 x座標の差を $dx$ ，y座標の差を $dy$ で表しています。
$dx=0$ のときは， $dx$ に何をかけても0なので，考慮する必要はありません。
$dx=i$ の選び方とき，2点PQの選び方は次の3つの自由度から定められるので， $(N-i) M^{2}$ 通りになります。
この計算は，ループ処理で1からN-1まで距離を求めればよいので，計算量 $O(N)$ で求まります。

1つ目の点の $x$ 座標を 1 から $N-x$ の中からどれにするか
行の選び方（クリックすると展開されます）
ここの"行の選び方"は， $1 \leqq x_i \leqq x_j \leqq N$ のうち， $x_j - x_i = dx$ となる $(x_i,\ x_j)$ の選び方です。
例えば， $x=(1, 2, 3)$ のうち，差が 1 となる選び方は， $(x_i,\ x_j) = (1, 2), (2, 3)$ の2通りとなります。
1つ目の点の $y$ 座標を 1から $M$ の中からどれにするか
2つ目の点の $y$ 座標を 1から $M$ の中からどれにするか

y座標についても，x座標と同様に計算ができます。最後に，STEP 2 で求めた残りの点の場合の数を考慮して， $x$ の寄与と $y$ の寄与の総和を求めると答えになります。

$\displaystyle ans = \biggl( \sum_{i=1}^{N-1} {i(N-i)M^{2}} + \sum_{i=1}^{M-1} {i(M-i)N^{2}} \biggr) \times {}_{NM-2} C_{K-2}$

コード

提出コード: リンク
MODnCr: GitHub

def main():
    MOD = 1_000_000_007
    N, M, K = map(int, input().split())
    
    ans= 0
    # axis: x
    for d in range(1, N):
        ans += d * (N-d)*M**2
        ans %= MOD
    # axis: y
    for d in range(1, M):
        ans += d * (M-d)*N**2
        ans %= MOD
 
    nCr = MODnCr(mod=MOD)
    ans *= nCr(N*M-2, K-2)
    ans %= MOD
    print(ans)
    return

解説つき！競ってわかる　はじめてpythonプログラミング

競プロの情報発信と備忘録を兼ねている

問題

解説

具体例 $(N, M, K) = (2, 2, 3)$ で考える

STEP 1

STEP 2

STEP 3

コード

関連リンク

問題

解説

具体例 で考える

STEP 1

STEP 2

STEP 3

コード

関連リンク

具体例 $(N, M, K) = (2, 2, 3)$ で考える