AOJ - コードをもっと短くしたい

Question

AOJの
Matrix Vector Multiplication
からです。行列と配列の内積をnumpyを使用せずに求める問題で
下記コードにてacceptedしたのですが、
個人的に改善したい点があります。
➀range(len(M)),range(len(ans))←コードの視認性を悪くする
➁vec = 0←他の簡潔な方法があるのではないか
➂そもそものコードが長いような気がする(forを多用している気がする、関数にして呼び出すのもありだが、そもそもの行数をもっと短くしたい)
以上の点についてアドバイスをお願いできないでしょうか。

 
n, m = map(int, input().split())

M = []
for i in range(1, n+1):
    row = list(map(int, input().split()))
    M.append(row)

v = []
for i in range(1, m+1):
    v.append(int(input()))


ans = []

for i in range(len(M)):　　　　　　　　　　　　　　　　　←
    vec = 0　　　                                     ←
    for j in range(len(M[0])):
        b = M[i][j] * v[j]
        vec += b
    ans.append(vec)

for i in range(len(ans)):　　                         ←
    print(ans[i])

Answer 1

だいぶすっきりしました。

def dot_vec(vec1, vec2):
    # assert len(vec1) == len(vec2)   競プロなら不要か 
    return sum(
        e1 * e2 for e1, e2 in zip(vec1, vec2)
    )

ans = [
    dot_vec(row, v) for row in M
]

• 関数を適宜定義する
• 内包表記やジェネレータ式、組み込み関数を活用する

---

出力ももっとすっきり書けます。

print(*ans, sep='\n')

あるいは

for e in ans:
    print(e)

range(len(...)) を使わなければいけない場面はそんなに多くない筈です。

Comments

[esklia]

ご回答くださりありがとうございます。内包表記とzipでだいぶ変わりますね、、出力も工夫の余地がまだまだあることが分かりました。ありがとうございました。

Answer 2

M = [
    [10, 20, 30],
    [40, 50, 60]
]
v = [1, 2, 3]

ans = [sum(e * v[i] for i, e in enumerate(row)) for row in M]
print('\n'.join(map(str, ans)))

参考情報

• Pythonで配列の中身をprintするときの実行速度比較

https://qiita.com/chiruno/items/21b65573bb696f4ce3ce

Comments

[esklia]

ご回答くださりありがとうございます。printの仕方で実行速度が異なるというのはあまり考えていませんでしたのでとても勉強になります。記事の方も拝見させていただきます。有益な記事と実装例を教えて下さりありがとうございました。