AtCoder「競プロ典型90問」の004-CrossSumのTLEを解消したいです。

Question

前提・実現したいこと

お世話になります。
AtCoderの「競プロ典型90問」の004-CrossSumについて質問がございます。
Python(3.8.2)にて以下のコードで提出したところ、2つのテストケースにおいてTLEとなってしまいます。
改善点などございましたら、どなたかご教授いただけますと幸いです。
宜しくお願い致します。

該当のソースコード

# 004-Cross Sum
import numpy as np
H,W=map(int, input().split())
A=np.array([list(map(int, input().split())) for i in range(H)])
Acolumn=np.sum(A, axis=0)
Arow=np.sum(A, axis=1)
B=[]
for i in range(H):
    b=[]
    for j in range(W):
        b.append(Arow[i]+Acolumn[j]-A[i,j])
    B.append(b)
for b in B:
    print(*b)

Answer 1

こちら気になったので色々調べてみたのですが、
結果としてはNumpyをインストールしていることでTLEになっているようでした。

以下調査した内容を記載します。

元にしたACコード

 
H, W = map(int, input().split())
A = [[int(d) for d in input().split()] for _ in range(H)]
rows = [sum(a) for a in A]
cols = [sum(z) for z in zip(*A)]
 
 
for y, row in enumerate(A):
	print(' '.join(str(rows[y] + cols[x] - A[y][x]) for x, v in enumerate(row)))

このコードはSimmah_9130様と考え方はほとんど変わりませんでしたので
すこしずつSimmah_9130様のコードに置換してみました。
すると以下のように

import numpy as np

だけを記載して回答した時にTLEとなり原因が

import numpy as np

だとわかりました。

import numpy as np
 
H, W = map(int, input().split())
A = [[int(d) for d in input().split()] for _ in range(H)]
rows = [sum(a) for a in A]
cols = [sum(z) for z in zip(*A)]
 
 
for y, row in enumerate(A):
	print(' '.join(str(rows[y] + cols[x] - A[y][x]) for x, v in enumerate(row)))

実行時間を確認してみたのですが、
numpyをインポートするコードを残している方だけ約100msだけ長く処理をしており
ギリギリ2000msを超えてしまいTLEとなっていたようです。

numpyは便利ですが地味にこういった落とし穴があるのがわかり非常にありがたい質問でした。
ありがとうございました。

Comments

[Simmah_9130]

Tony1r0n様
早速の対応、ありがとうございました。
numpyを用いることで処理時間が短くなるかと思いましたが、
まさかそれが仇になるとは思ってもみなかったです。
代替のコードまでご提示いただき、大変参考になりました。
またご縁がございましたら、その際は宜しくお願い致します。