[Python] ABC018 C 100点 いもす法

問題

C – 菱型カウント

回答

スライドの解説に従うとTLEになってしまいました。

以下を参考にして、いもす法を用いました。

ABC018 C – 菱型カウント

# https://emtubasa.hateblo.jp/entry/2018/11/09/031048 参照

R, C, K = map(int, input().split())
S = [input() for _ in range(R)]

# imos[x][y] := (x, y)を中心として菱形を作成した時、
# その範囲内にある黒マスの個数
imos = [[0] * (C+1) for _ in range(R)]

for r in range(R):
    for c in range(C):
        if S[r][c] == 'x':
            for k in range(K):
                if r-k >= 0:
                    imos[r-k][max(0, c-K+1+k)] += 1
                    imos[r-k][min(C-1, c+K-k)] -= 1
                if r+k < R:
                    imos[r+k][max(0, c-K+1+k)] += 1
                    imos[r+k][min(C-1, c+K-k)] -= 1

for r in range(R):
    for c in range(1, C+1):
        imos[r][c] += imos[r][c-1]

ans = 0
# 中心 (r, c) を固定 
for r in range(K-1, R-K+1):
    for c in range(K-1, C-K+1):
        if imos[r][c] == 0:
            ans += 1

print(ans)