AtCoder ABC127 Eを寄与で解く：距離総和の式と組合せ前計算、ABC330 Dへの応用

全ての配置を列挙して距離を足す前に、距離の定義が座標ごとの寄与の和へ分けられるかを確認します。分けられるなら、差 d を持つ2マスの組数を数え、固定した2マス以外の選び方を組合せで掛けます。ABC127 E「Cell Distance」では、V = N × M と置き、横方向（N 個並ぶ座標）と縦方向（M 個並ぶ座標）を別々に足せば、全配置のマンハッタン距離の総和を求められます。

この記事の式では、盤面の描き方で横・縦が逆に見えることを避けるため、x を N 側の座標、y を M 側の座標として固定します。距離がこの二つの差の和になっていること、固定したペアを含む配置の残りを一様に選べることが、この数え上げを使える条件です。条件が変わるABC330 Dでは、距離の差ではなく、固定した o マスから同じ行と列にある別の o マスを選ぶ積へ考え方を移します。

1次元の距離和を差ごとに追う

以下は公式サンプルではない、説明用の独自検算です。まず、位置1、2、3、4から異なる2点を選び、距離の総和を直接確かめます。同じ差を持つ組をまとめると、差が大きいほど組数は減りますが、1組あたりの寄与は大きくなります。

差 d	該当する組	組数	差の合計
1	(1,2), (2,3), (3,4)	3	3 × 1 = 3
2	(1,3), (2,4)	2	2 × 2 = 4
3	(1,4)	1	1 × 3 = 3

長さ L の1次元列で差を d に固定すると、左端の位置を選ぶ方法が L - d 通りあります。したがって差の総和は、差ごとの寄与を足す Σ d(L - d) です。上の例では 3 + 4 + 3 = 10 となり、閉じた形の L(L - 1)(L + 1) / 6 に L = 4 を入れても 4 × 3 × 5 / 6 = 10 になります。この「差」と「その差を持つ組数」を分ける見方が、2次元の各座標へそのまま移ります。

ABC127 Eはマンハッタン距離を二つの寄与に分ける

ABC127 Eの公式問題文では、N 行 M 列のマス目から K マスを選び、選んだマスを (xi, yi) としたときのコストを、全ての i < j について |xi - xj| + |yi - yj| を足した値として定義しています。求めるのは、取り得る全てのマス集合のコストの総和を 109 + 7 で割った余りです。公式制約は 2 ≤ N × M ≤ 2 × 105、2 ≤ K ≤ N × M です。

ABC127 Eの公式解説では、距離の中にある絶対値の和は、x 方向の差と y 方向の差を別々に集計して最後に足せます。ここで V = N × M と置き、まず距離へ寄与する異なる2マスを固定します。その2マスを含む配置を完成させるには、残りの V - 2 マスから K - 2 マスを選べばよいので、固定ペア1組が現れる配置数は W = C(V - 2, K - 2) です。C(a, b) は a 個から b 個を選ぶ組合せの数を表します。

横方向（N側の座標差）の寄与

x 座標の差を d とします。差が d になる座標の組は、N - d 通りです。その各座標に対して、最初のマスの列を M 通り、2番目のマスの列を M 通り選べるため、セルのペアは (N - d) × M2 組あります。各ペアは距離へ d を加えるので、横方向の寄与は次の式です。

Sx = M2 × Σd = 1..N - 1 d(N - d)

ここで N - d は差 d を持つ行（N 側の座標）の組数、M2 はそれぞれの行から列を選ぶ組数、d は1組あたりの x 方向の距離です。1次元の閉形式を使うと、Sx = M2 × N(N - 1)(N + 1) / 6 になります。

縦方向（M側の座標差）の寄与

y 座標でも同じ数え方をします。差 d の列の組は M - d 通りあり、各列から行を選ぶ方法が N2 通りあります。したがって、縦方向の寄与は Sy = N2 × Σd = 1..M - 1 d(M - d) で、閉形式は Sy = N2 × M(M - 1)(M + 1) / 6 です。M - d は列の組数、N2 は各列から行を選ぶ組数、d は y 方向の差を表します。

固定ペア以外の選び方 W は、横方向でも縦方向でも同じです。よって答えは、二つの軸の寄与にこの係数を掛けて足した answer = W × (Sx + Sy) mod (109 + 7) です。x と y の両方が異なるペアは、マンハッタン距離の中でそれぞれの差を一度ずつ持つため、二つの和に分けても不足や重複はありません。

固定ペアに組合せを掛けてよい理由

全配置を一つずつ作る代わりに、まず「この2マスの x 方向の差が、全配置のコストへ何回現れるか」を考えます。異なる2マスを固定したとき、そのペアを含む配置は、残りの V - 2 マスから K - 2 マスを選ぶことで過不足なく作れます。したがって、固定ペアの差は W 個の配置で1回ずつ現れます。全ての固定ペアについてこの寄与を足すと、各配置の各ペアの距離をちょうど一度ずつ足したことになります。

この数え方は、距離が座標ごとの和に分かれ、固定ペアを選んだ後の残りの選択数がペアの位置に依存しない場合に使えます。距離が座標の最大値など別の形で定義されている場合や、残りのマスに追加条件がある場合は、同じ Sx + Sy や C(V - 2, K - 2) をそのまま当てはめず、固定した対象の寄与と残りの選び方を改めて数えます。

2×2盤面を6配置すべて列挙する

N = 2、M = 2、K = 2 では、4マスから2マスを選ぶので配置は6通りです。次の表は公式問題文から計算した独自の全列挙です。入力と出力が公式サンプル1と同じであることと、表自体が公式サンプルの転載ではないことを分けて確認します。

選んだ2マス	|Δx|	|Δy|	距離
(1,1), (1,2)	0	1	1
(1,1), (2,1)	1	0	1
(1,1), (2,2)	1	1	2
(1,2), (2,1)	1	1	2
(1,2), (2,2)	1	0	1
(2,1), (2,2)	0	1	1

距離の合計は 1 + 1 + 2 + 2 + 1 + 1 = 8 です。式でも、V = 4、W = C(2, 0) = 1 となり、Sx = (2 - 1) × 22 × 1 = 4、Sy = 4 なので、答えは 1 × (4 + 4) = 8 です。全列挙で見た6つの距離が、横と縦の寄与へ分けた計算でも同じ合計になります。

ABC127 Eの公式サンプル3件

次の入出力は、公式問題文に掲載されたサンプルを加工せずに示したものです。後の参考コードへ、各入力をそのまま与えて確認します。

サンプル1

2 2 2

期待出力：

8

サンプル2

4 5 4

期待出力：

87210

サンプル3

100 100 5000

期待出力：

817260251

組合せを前計算する完全Python参考コード

W = C(V - 2, K - 2) を毎回の組合せ計算で求めると、階乗の計算を繰り返すことになります。そこで fact[i] = i! と、その逆元階乗 inv_fact[i] = (i!)-1 を V まで前計算します。fact[V] の逆元を一度求め、そこから inv_fact[i - 1] = inv_fact[i] × i と逆向きに作れば、全ての値を線形時間で用意できます。

公式制約では V ≤ 2 × 105 < 109 + 7 なので、MOD = 109 + 7 は素数であり、V 以下の階乗は MOD の倍数になりません。Fermatの小定理により、a の逆元は aMOD - 2 mod MOD で求められます。コードでは各乗算の後に % MOD を取り、軸の閉形式は整数として6で割ってから最後に剰余を取ります。

以下は問題文と公式解説を基に、この記事用に組み立てた参考実装です。AtCoder公式コード・提出物ではなく、AC結果や実測性能を示すものでもありません。

import sys

MOD = 10**9 + 7


def main():
    n, m, k = map(int, input().split())
    v = n * m

    fact = [1] * (v + 1)
    for i in range(1, v + 1):
        fact[i] = fact[i - 1] * i % MOD

    inv_fact = [1] * (v + 1)
    inv_fact[v] = pow(fact[v], MOD - 2, MOD)
    for i in range(v, 0, -1):
        inv_fact[i - 1] = inv_fact[i] * i % MOD

    def comb(a, b):
        if b < 0 or b > a:
            return 0
        return fact[a] * inv_fact[b] % MOD * inv_fact[a - b] % MOD

    ways = comb(v - 2, k - 2)
    sx = m * m * n * (n - 1) * (n + 1) // 6
    sy = n * n * m * (m - 1) * (m + 1) // 6
    answer = ways * ((sx + sy) % MOD) % MOD
    print(answer)


if __name__ == '__main__':
    main()

入力は公式どおり N M K の1行です。v が盤面全体のマス数、ways が固定ペアの残りを選ぶ W、sx と sy が上で導いた二つの軸の寄与に対応します。comb の境界分岐は、b が範囲外なら0を返すためのものです。公式制約内では a = v - 2、b = k - 2 は有効な範囲に入ります。

正しさ、mod、境界、計算量を確認する

1. 距離の分解：各ペアの距離は |Δx| + |Δy| なので、全配置の総和も x 方向の総和と y 方向の総和の和になります。
2. 固定ペアの回数：異なる2マスを固定すると、残りは V - 2 マスから K - 2 マスを選ぶ W 通りです。各配置の各ペアは、固定ペアとして一度だけ数えられます。
3. 差ごとの組数：N 側で差 d の座標の組は N - d 通り、列の選び方は M2 通りです。M 側も同様に入れ替えればよく、1次元の和を閉形式へ置き換えています。

K = 2 なら W = C(V - 2, 0) = 1 で、選んだペア以外を選ぶ必要がありません。K = V でも W = C(V - 2, V - 2) = 1 です。N = 1 または M = 1 のように一方の軸の長さが1なら、その軸の式にある N - 1 または M - 1 が0になり、もう一方の軸だけが残ります。公式制約では V ≥ 2 なので、コードは v - 2 以下の階乗を参照できます。

階乗と逆元階乗の前計算は O(V) 時間・O(V) メモリです。閉形式による sx、sy と答えの計算は O(1) 時間なので、コード全体も O(V) 時間・O(V) メモリです。これはループと配列から分かる計算量であり、特定環境での実測時間を表すものではありません。

このコードへABC127 Eの公式サンプル3件を無加工で入力すると、期待出力は順に 8、87210、817260251 です。2×2の全列挙も同じコードの式で 8 になり、公式サンプル確認と独自の境界確認を別々に追えます。

ABC330 Dでは固定点から行と列の積へ移す

ABC330 Dの公式問題文では、N × N の盤面にある o と x のうち、相異なる3マスを選びます。3マスが全て o で、ちょうど2マスが同じ行、ちょうど2マスが同じ列になる三つ組の個数を求めます。公式制約は 2 ≤ N ≤ 2000 です。N2 マスから3マスを直接選んで条件を調べると、候補数は C(N2, 3) で、O(N6) まで増えます。そこで三つ組を直接列挙せず、三つ組の行と列が交わる o マスを1つ固定します。各行・各列の o の数を前計算し、全 N2 マスを走査すれば O(N2) で数えられます。ABC127 Eのように距離を足す条件ではないため、d(N - d) の和や C(V - 2, K - 2) は使いません。

ABC330 Dの公式解説が示す移し替えは、三つ組の行と列が交わる o マスをアンカーとして固定することです。アンカーを (i, j) とすると、同じ行にある別の o マスは row[i] - 1 通り、同じ列にある別の o マスは col[j] - 1 通りです。アンカー自身を1回ずつ引くため、アンカーの寄与は (row[i] - 1) × (col[j] - 1) となります。

同じ行と同じ列の条件を同時に満たすマスは、この三つ組ではアンカー1つに決まります。そのため、全ての o マスをアンカーにしてこの積を足せば、同じ三つ組を別のアンカーから重ねて数えません。ABC127 Eで「固定ペアの後を選ぶ」だった部分が、ABC330 Dでは「固定点から行と列の候補を独立に選ぶ」へ変わった、と見ると両問のつながりを追えます。

式で書くと、答えは ΣS[i][j] = 'o' (row[i] - 1) × (col[j] - 1) です。コードの i と j は0始まりですが、row[i] はその行の o の数、col[j] はその列の o の数という意味は説明中の1始まりの座標と同じです。

公式サンプル1をアンカーで追う

3
ooo
oxx
xxo

この盤面の行ごとの o の数は row = [3, 1, 1]、列ごとは col = [2, 1, 2] です。アンカー (1,1) の寄与は (3 - 1) × (2 - 1) = 2、アンカー (1,3) の寄与も 2 です。他の o マスは同じ行または同じ列に別の o がないため寄与0となり、合計は公式出力の 4 になります。この追跡は、公式の入力と期待出力を使って固定点の積を確かめた、この記事独自の検算です。公式解説そのものや、その追加資料ではありません。

ABC330 Dの完全Python参考コード

コードは、まず盤面を読み、行と列の o の数を集計します。その後、もう一度全マスを走査し、o マスだけをアンカーとして積を足します。入力は公式どおり、最初に N、続けて長さ N の文字列を N 行読みます。

以下は問題文と公式解説を基に、この記事用に組み立てた参考実装です。AtCoder公式コード・提出物ではなく、AC結果や実測性能を示すものでもありません。

import sys


def main():
    input = sys.stdin.readline
    n = int(input())
    grid = [input().strip() for _ in range(n)]

    row = [0] * n
    col = [0] * n
    for i in range(n):
        for j in range(n):
            if grid[i][j] == 'o':
                row[i] += 1
                col[j] += 1

    answer = 0
    for i in range(n):
        for j in range(n):
            if grid[i][j] == 'o':
                answer += (row[i] - 1) * (col[j] - 1)
    print(answer)


if __name__ == '__main__':
    main()

行と列の集計、およびアンカーの走査は、それぞれ N × N のマスを高々一度ずつ調べるため、コードから読める計算量は O(N2) 時間です。盤面 grid と行・列の配列を保持するので、追加メモリはこの実装では O(N2) です。ここでも、これは実装上の計算量であり、実測時間や提出結果を示す主張ではありません。

ABC330 Dの公式サンプル3件

次の3件も、公式問題文の入力と期待出力を加工せずに示します。サンプル1のアンカー追跡と、参考コードへ標準入力を渡す確認を区別します。

サンプル1

3
ooo
oxx
xxo

期待出力：

4

サンプル2

4
oxxx
xoxx
xxox
xxxo

期待出力：

0

この盤面では、各 o マスについて、その行の o も列の o も1個だけです。どのアンカーでも row[i] = 1、col[j] = 1 なので、寄与は (1 - 1) × (1 - 1) = 0 となり、合計も0です。

サンプル3

15
xooxxooooxxxoox
oxxoxoxxxoxoxxo
oxxoxoxxxoxoxxx
ooooxooooxxoxxx
oxxoxoxxxoxoxxx
oxxoxoxxxoxoxxo
oxxoxooooxxxoox
xxxxxxxxxxxxxxx
xooxxxooxxxooox
oxxoxoxxoxoxxxo
xxxoxxxxoxoxxoo
xooxxxooxxoxoxo
xxxoxxxxoxooxxo
oxxoxoxxoxoxxxo
xooxxxooxxxooox

期待出力：

2960

ABC127 Eでは、座標差を足せることを確認してから、差ごとのペア数と残りの選択数を掛けます。ABC330 Dでは、距離の和という条件を捨て、三つ組を一意に決めるアンカーを固定して、行と列の残り個数を掛けます。次の問題文を読むときは、全列挙を始める前に「何を固定すれば、残りの選び方を同じ係数または独立した積として再利用できるか」を確認し、公式サンプルをそのまま実行して式とコードの対応を確かめてください。