【備忘録】競プロ典型90問 004(★2) #1

　競プロ典型90問では、各問題に難易度（★の数）が設定されています。★1の問題はないようでしたので、まずは★2から解いていこうと思います。

　備忘録ということで、後から見返したときに分かりやすいように

• 問題
• アイデア
• 改良案を考える
• 最終的な実装例

の流れでまとめていこうと思います。
　自力でできることがなくなったら解説を参照していきます。

　僕のような初心者の方に主に多く見ていただいて切磋琢磨したいと考えていますが、レベル問わず皆さまの閲覧やコメントをお待ちしています。

問題

　行列$A$の各要素$A_{i, j}$に対して、その要素が含まれる$i$行と$j$列のすべての要素の和 $$B_{i, j} = A_{i, j} + \sum_{i \neq j} A_{i, j} + \sum_{j \neq i} A_{i, j}$$を求める。

アイデア

アルゴリズム

　2つのforループを用いて、$B_{i, j}$を愚直に計算するのが最も直感的（でも経験上、TLEなりがち）。
…このときの計算量は、2つのforループで$O(H\times W)$、行と列に対する和でそれぞれ$O(W)$、$O(H)$の計$$O(H \times W) \times (O(W) + O(H))$$となり、$W, H \approx 10^3$のもとでは$10^9$のオーダーになってしまう。。。
⇒このままではいけない！

その他

　行列形式の入力の受け取り方と出力の仕方をどうしよう。普段は

# 入力
A = []
for _ in range(H):
   tmp = list(map(int, input().split()))
   A.append(tmp)

# 出力（Bを出力したい）
B = [[0] * W for i in range(H)]
for i in range(H):
   res = ""
   for j in range(W):
       res += str(B[i][j])
       if j < W - 1:
           res += ' '
   print(res)

という感じ。入力にmap関数が便利なのはグプタ（GPT）に教えてもらった。🥰
スペース区切りでの出力を要求されると、もっと簡潔に書けるはずだよなぁと思いながらも、いつもこんな感じで書いてます。。。

改良案を考える

アルゴリズム

　競技プログラミング初心者として、ABCのC問題に挑戦する中で僕が何となく感じていたことが、この問題にも共通している気がします。それは「問題の手順通りに素直に計算してはいけない」ということです。

　問題はあくまでも、どのような入力があってどのような出力をするべきかを、人間の僕たちにとって分かりやすい形に書かれた計算手順を添えて教えてくれるものです。ループや条件分岐ができれば、手順をそのままPythonに翻訳することはできますが、経験則上、計算速度には概して向上の余地があるようです。

　競技プログラミングの観点から僕たちがするべきことは、日本語からPythonへの言わば直訳ではなく、

• 問題を理解する（どんな入力からどんな出力か）
• 計算速度の上で有利な手順を改めて考える（重複はないか、愚直すぎないか）
• 最速と思われる手順をPythonに翻訳する

ことだと最近気づきました。

　その上で今回の問題を再考すると、各$B_{i, j}$を求めるたびに行と列の和をとっている！
⇒$B_{i, 1}, B_{i, 2}, \dots, B_{i, W}$を求めるには$\sum_j A_{i, j}$の計算が必要だが、毎回計算し直すと重複してしまう。。ということで改善案：

• 行列Aを入力から受け取り、すぐに各行と各列の和を計算した配列をつくる
• その配列を参照して$B$の各要素の計算をする

計算量は$O(H\times W) + O(H\times W) + O(H\times W) \approx O(H\times W)$となります。

出力

　少し調べただけで、これまでいかに自分が冗長なコードを書いていたかということが分かりました。リストを空白区切りで出力する方法には、次の2つがいいようです。

• joinを用いる

B = [[0] * W for i in range(H)]
for i in range(H):
    print(' '.join(map, B[i]))

• *（アンパック演算子）を用いる

B = [[0] * W for i in range(H)]
for i in range(H):
    print(*B[i])

　なんとなくアンパック演算子を用いるほうが簡潔に見えるので、一旦こっちを使ってみようと思います。

最終的な実装例

H, W = map(int, input().split())
A = []
for _ in range(H):
    tmp = list(map(int, input().split()))
    A.append(tmp)

# ここで各行の和rowsumと各列の和colsumを計算しておく
rowsum = []
for i in range(H):
    tmp = 0
    for j in range(W):
        tmp += A[i][j]
    rowsum.append(tmp)  
colsum = []
for j in range(W):
    tmp = 0
    for i in range(H):
        tmp += A[i][j]
    colsum.append(tmp)
    
B = [[0] * W for i in range(H)]
for i in range(H):
    for j in range(W):
        B[i][j] = rowsum[i] + colsum[j] - A[i][j]

for b in B:
    print(*b)

　drkenさんが公開されている解説を見たところ、その行数の短さに度肝を抜かれました。

まだまだ改良の余地はありそうですが、この記事はここで終わって余裕があれば続編を書きます！お読みいただきありがとうございました。