【競プロ典型90問】004をPythonで解説

はじめに

最近競プロを始めたので，競プロ典型90問の★1~★3を解いて典型問題へのアプローチ方法を学んでいます．
公式のソースコードがC++で困っているPython勢の方も多いと思うので，自分の解答をシェアしたいと思います！
初心者の解答なので，もっといい方法があれば教えていただけると嬉しいです！

004 Cross Sum（★2）

問題文

$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,W≤2000$
• $1≤A_{i,j}≤99$
• 入力は全て整数

解法

各マス $(i,j)$ について第 $i$ 行と第 $j$ 列の整数の合計を求めていくだけなので，やること自体は簡単です．しかし，愚直に実装しようとすると，第 $i$ 行の合計を求めるために $O(W)$，第$j$ 列の合計を求めるために $O(H)$ の計算量がそれぞれかかってしまい，これをすべてのマスについて計算するので $O((H+W)HW)$ の計算量になり TLE　してしまいます．
そこで，各マスについての個別の計算は必ず必要だと思われるので，各行と各列の合計を先に求めておくことを考えます．
まず，各行の合計を記録するためのリスト sum_row と 各列の合計を記録するためのリスト sum_col を用意し，第 $i$ 行の合計を sum_row[i] に，第 $j$ 行の合計を sum_col[j] に入れることにします．
このとき，「マス $(i,j)$ と同じ行または同じ列にあるマス（自分自身を含む）に書かれている整数をすべて合計した値」を ans[i][j] とすると以下の式で表されます．

ans[i][j] = sum\_row[i] + sum\_col[j] - A[i][j]

ただし，sum_row[i] ，sum_col[j] にはそれぞれ A[i][j] が1回ずつ足されているため，重複している分を引いています．
sum_row を作るための計算量を考えると，1行の合計を求めるための計算量は $O(W)$ で，それを $H$ 行分計算するので，$O(HW)$．sum_col についても同様です．
したがって，各マスについての計算と，行，列ごとの合計の計算を分離したため，全体の計算量が $O(HW)$ になりました．制約は $2≤H, W≤2000$　なので，十分間に合いそうです.

実装

解法の部分で説明した内容をそのままコードにすればよいので，それほど難しいところはないと思いますが，sum_col の計算が複雑になってしまう人がいるかもしれません．そこで，今回は map を使うことで1行で実装できるようにしています．少し見慣れない書き方かもしれませんが，2次元配列の列ごとの合計を求める際に非常に便利なので，覚えておくとよいと思います．

↓↓↓ 2次元リストの合計の求め方です ↓↓↓

answer.py

import sys
input = sys.stdin.readline


H, W = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(H)]
# 行, 列ごとの合計をリストに入れる
sum_row = list(map(sum, A))
sum_col = list(map(sum, zip(*A)))
# Aのすべての要素を更新
for i in range(H):
    for j in range(W):
        A[i][j] = sum_row[i] + sum_col[j] - A[i][j]
# 出力
for i in range(H):
    print(*A[i])

さいごに

今回はやること自体はそれほど難しくないものの，計算量を減らす工夫が必要な問題でした．同じ計算や値を何度も使う場合には，それらをはじめに求めておくと上手くいくことが多いと思います．
よく分からないことやコードの改善点などありましたら，気軽にコメントしてください．

最後までお付き合いいただきありがとうございました！