ABC278-E 一部のテストケースがTLEになってしまう
解決したいこと
AtCoderのABC278-E問題(https://atcoder.jp/contests/abc278/tasks/abc278_e)
で、一部のテストケース(24個中3個)のTLEが取れません。改善方法を教えていただけますでしょうか。言語はPyPyを使用しています。
発生している問題・エラー
03_max_13.txt,04_edge_16.txt,04_edge_20.txtがTLEになる。
該当するソースコード
from collections import defaultdict
'''二次元累積和作成'''
def make_cumulative_sum_2D(A) -> list:
H,W = len(A),len(A[0])
S = [[0] * (W+1) for _ in range(H+1)]
for i in range(1,H+1):
for j in range(1,W+1):
S[i][j] = A[i-1][j-1]+S[i][j-1]+S[i-1][j]-S[i-1][j-1]
return S
'''2次元累積和Sの[h1,h2)x[w1,w2) の区間和を求める'''
def interval_sum_2D(S,h1,w1,h2,w2):
return S[h2][w2] - S[h1][w2] - S[h2][w1] + S[h1][w1]
def main():
H,W,N,h,w = map(int,input().split())
A = [list(map(lambda x:int(x)-1,input().split())) for _ in range(H)]
L = []
C = defaultdict(int)
i_list = [[[0] * W for _ in range(H)] for _ in range(N)]
for j in range(H):
for k in range(W):
i_list[A[j][k]][j][k] += 1
for i in range(N):
L.append(make_cumulative_sum_2D(i_list[i]))
for i in range(H):
for j in range(W):
C[A[i][j]] += 1
for i in range(H-h+1):
l = []
for j in range(W-w+1):
ans = 0
for k in range(N):
if C[k]-interval_sum_2D(L[k],i,j,i+h,j+w) >= 1:
ans += 1
l.append(ans)
print(*l)
if __name__ == "__main__":
main()
自分で試したこと
・グラフ作成回数の削減
・main関数への埋め込み
→どちらも改善せず
0