PythonでJOI難易度6を埋める #9

banner - 横断幕 (Banner)

愚直に計算すると組み合わせ数が爆発して死．
そこで，4個の頂点のなかで，1個の頂点の色は自由に選べることに注目する．長方形から1つの頂点を取り除いた形で，他の3個の頂点の色が全て異なるような図形を考える．その図形はのこりの1頂点に，どの色を当てはめても条件を満たす組合せとなるので，そのような図形の数を数える．事前にそれぞれの行，列で色の個数を数えておく．点が黒色だとすると，求めたい図形の個数は，同じ行の白の個数 同じ列の灰の個数 + 同じ行の灰の個数 同じ列の白の個数となる．これをそれぞれの全ての頂点で行ない，その合計をとる．今数えた図形から，まったく同じ2個の長方形ができるので，合計を2で割ったものが答え．図を用いた分かりやすい解説

from sys import stdin
input = stdin.readline

h, w = map(int, input().strip().split())
maps = [list(map(int, input().strip().split())) for _ in range(h)]

count_maps_h = [[0, 0, 0] for _ in range(h)] #黒，灰，白
count_maps_w = [[0, 0, 0] for _ in range(w)] #黒，灰，白

for i in range(h):
    for j in range(w):
        if maps[i][j] == 0:
            count_maps_h[i][0] += 1
            count_maps_w[j][0] += 1
        if maps[i][j] == 1:
            count_maps_h[i][1] += 1
            count_maps_w[j][1] += 1
        if maps[i][j] == 2:
            count_maps_h[i][2] += 1
            count_maps_w[j][2] += 1

ans = 0

for i in range(h):
    for j in range(w):
        if maps[i][j] == 0:
            ans += count_maps_h[i][1] * count_maps_w[j][2] + count_maps_h[i][2] * count_maps_w[j][1]
        if maps[i][j] == 1:
            ans += count_maps_h[i][0] * count_maps_w[j][2] + count_maps_h[i][2] * count_maps_w[j][0]
        if maps[i][j] == 2:
            ans += count_maps_h[i][0] * count_maps_w[j][1] + count_maps_h[i][1] * count_maps_w[j][0]

print(ans // 2)