Atcoder典型90問 004 - Cross Sum（★2）

▶︎https://atcoder.jp/contests/typical90/tasks/typical90_d

問題

H 行 W 列のマス目があります。上から i (1≤i≤H) 行目、左から j (1≤j≤W) 列目にあるマス
(i,j) には、整数 A(i,j)が書かれています。
すべてのマス (i,j) (1≤i≤H,1≤j≤W) について、以下の値を求めてください。

マス (i,j) と同じ行または同じ列にあるマス（自分自身を含む）に書かれている整数をすべて合計した値

制約

2≤H≤2000 \\
2≤W≤2000 \\
1≤A_{(i,j)}≤99 \\

試したこと

各マスに対して愚直に行/列の足し算を行うと、最低でもO(HW * [H + W])となりPythonには荷が重い計算量になる。
- あとでわかったことだが、input時以外にO(N^2)の計算が入った時点でほぼ詰んでしまう。
- その理由はPythonでは配列宣言がとくに重いから。
- https://tech-lab.sios.jp/archives/21127#i-10
上のアルゴリズムでは同じ行と同じ列の足し算を重複して行っていることに気づければ、あとは計算量を乗り越える部分に入る。とはいえO(HW)がギリギリなので、かなりアルゴリズムを気をつけないと定数倍さえ響く。かなり苦労した...
最初はnumpyで実装したが、Pypy3とnumpyが併用できないのでnumpyを諦めた。
- 結果的には、実行速度が早い順に
- Pypy3 & map > numpy > Pypy3&for
- となった(左が早い)。numpyがTLE２個、Pypy3&forがTLE4個くらい。

適切な方針

各行と各列の和を事前に計算しておく(前処理)
求める行列B[i][j]は、与えられた行列Aのi行目の和とj行目の和から重複部分(A[i][j]を一つ分)除いた値である。
行列の列和を求める際はmapを使うとよい。単に二重for文でゴリ押そうとするとTLEコース。
- 二重for文ではなく、逆行列を作ってその行和を取ろうとしたのがミスだったのかもしれない。
- いずれにせよmap文を使うのが大切そう。

解説付きコード

[H, W] = list(map(int, input().split()))
ans_mat = [[0 for i in range(W)] for j in range(H)]
tar_mat = [[0 for i in range(W)] for j in range(H)]

#行和と列和の保存用配列
addH_list = [0 for i in range(H)]
addW_list = [0 for i in range(W)]

#入力受け取り
for i in range(H):
  input_list = list(map(int, input().split()))
  tar_mat[i] = input_list.copy()

#行和の保存
for i in range(H):
  addH_list[i] = sum(tar_mat[i])

#列和の保存
for j in range(W):
  addW_list[j] = sum([tar_mat[i][j] for i in range(H)])

#行和と列和の更新
for i in range(H):
  for j in range(W):
    ans_mat[i][j] = addH_list[i] + addW_list[j] - tar_mat[i][j]

#配列をスペースで結合した文字列として出力する
for i in range(H):
  print(" ".join(list(map(str, ans_mat[i]))))

付録：TLEを起こしたコード(二重for文)、

# 転置行列を返す関数(戦犯)
def transpose(mat):
  H = len(mat)
  W = len(mat[0])
  new_mat = [[0 for i in range(H)] for j in range(W)]
  for i in range(H):
    for j in range(W):
      new_mat[j][i] = mat[i][j]
  return new_mat

def main():
    [H, W] = list(map(int, input().split()))
    ans_mat = [[0 for i in range(W)] for j in range(H)]
    tar_mat = [[0 for i in range(W)] for j in range(H)]
     
    addH_list = [0 for i in range(H)]
    addW_list = [0 for i in range(W)]
     
    for i in range(H):
      input_list = list(map(int, input().split()))
      tar_mat[i] = input_list.copy()
     
    for i in range(H):
      addH_list[i] = sum(tar_mat[i])
     
    for j in range(W):
    　　　　# sumは組み込み関数なので早いだろうし転置行列に対して使えばいいだろう
    　　　　# と思ったが、見通しが甘かった。
      tar_mat_t = transpose(tar_mat)
      addW_list[j] = sum(tar_mat_t[j])
        
    for i in range(H):
      for j in range(W):
        ans_mat[i][j] = addH_list[i] + addW_list[j] - tar_mat[i][j]
     
    for i in range(H):
      anslist = list(map(int, ans_mat[i]))
      for j in range(W):
        if j!= W-1:
          print(anslist[j], end=' ')
        else:
          print(anslist[j])
main()