Atcoder 典型90問004-★2[前処理][配列処理]

■誤解答

方針自体はすぐに思いついた．求めたい場所に対して縦と横の総和からその場所の値を引けば良いと考え実装．この際，行列の縦と横の総和の取り方が分からず迷う．縦の総和はlistで行けるが，横の総和はnpにしないとsum関数を使って上手くできない．

import numpy as np
H, W = map(int, input().split())
ans = 0
a = [list(map(int, input().split())) for l in range(H)]
a = np.array(a)
ans = [[0 for i in range(W)] for j in range(H)]
ans = np.array(ans)
for i in range(H):
 for j in range(W):
   ans[i,j] = sum(a[i]) + sum(a[:,j]) - a[i,j]
   
for i in ans:
 print(*i)

上記のやり方は，愚直に各値を求めている．ループ自体はH*Wだから計算量はO(HW)かと思ったのがセンスが無い．実際は総和を求めているので，その処理にO(H+W)掛かっている．そのためTLE.

■改善案

　縦と横の総和は何度も同じ値を計算しているのがもったいないと気づいた．そのため予め用意しておけば良い．この際，最初はnp配列のままsum関数を使って処理していたが，おそらくnpを使うとpypyが使えずずっとTLEだった．そこでsum関数を使わず愚直に縦と横の総和をlistで求めるように変更し，pypyで通したらACできた．

from sys import stdin
H, W = map(int, stdin.readline().split())

a = [list(map(int, input().split())) for l in range(H)]

ans = [[0] * W for _ in range(H)]


#tate = [sum(a[:,i]) for i in range (W)]
#yoko = [sum(a[i]) for i in range(H)]
tate = [0] * W
for i in range(W):
 cnt = 0
 for j in range(H):
   cnt += a[j][i]
 tate[i] = cnt

yoko = [0] * H
for i in range(H):
 cnt = 0
 for j in range(W):
   cnt += a[i][j]
 yoko[i] = cnt

for i in range(H):
 for j in range(W):
   ans[i][j] = yoko[i] +  tate[j] - a[i][j]
   
 print(*ans[i])