aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 127 E - Cell Distance (python)

自力AC。差の絶対値の総和を求める。
こういうときは | n - m | = max(n,m) - min(n,m) なので、a より小さいものの数、大きいものの数を x 座標と y 座標で別々に求めたらよい。これの一次元版を解いたことある気がする。
1 <= x <= n, 1<= y <= m とすると x =max(x,i) となる i の数は合計で (x-1)*m 個ある
そして X座標が x となる頂点の数も m 個あるので、X座標が x で発生する絶対値の総和の変化は

  1. (x-1)*m*m - (n-x)*m*m となる。

あとはこれを x と y について調べる。
しかしこの総和は、全ての n*m 個の頂点から 2 個選んだときの総和である。
(つまり k=2 の時の答えになっている)
そして、n*m 個の頂点から k 個選んだ時に、特定の 2 頂点が含まれる組み合わせは、
n*m-2Ck-2 であるので、これをかけたものが答えとなる。
(n*m-2Ck-2 にたどり着くのが一番大変だった。先に2個決めた状態で残りの n*m-2 個の頂点から k-2 個を選ぶと考えるとすぐ分かる)
nCr のライブラリーはモジュラ係数のやつを貼るだけ

n,m,k = map(int,input().split())
mod = 10**9+7

def make_array_for_comb(N, mod=10**9+7):
    fact = [1,1]
    fact_inv = [1,1]
    inv = [0,1]
    for i in range(2, N+1):
        fact.append((fact[-1]*i) % mod)
        inv.append((-inv[mod%i] * (mod//i)) % mod)
        fact_inv.append((fact_inv[-1]*inv[i]) % mod)
    return fact, fact_inv

def comb(n, r, mod=10**9+7):
  if (r < 0) or (n < r):
      return 0
  r = min(r, n - r)
  return fact[n] * fact_inv[r] * fact_inv[n-r] % mod

fact, fact_inv = make_array_for_comb(n*m,mod)

l = comb(n*m-2,k-2,mod)
ans = 0
for a in range(1,n+1):
  ans += a*m*(a-1)*m
  ans -= a*m*(n-a)*m
ans %= mod

for b in range(1,m+1):
  ans += b*n*(b-1)*n
  ans -= b*n*(m-b)*n
ans %= mod

print(ans*l%mod)