【Python】多次元配列に値を代入するときの注意点

やりたかったこと

動的計画法に関する問題で，初期化した2次元配列(dp)に対して，ある数字(60)を代入しようとしていた．

[問題] Atcoder EDPC C-Vacation(https://atcoder.jp/contests/dp/tasks/dp_c)

以下のようなコードを書いていた．

dp=[[0,0,0]]*5
print(dp)
#出力 [[0,0,0],[0,0,0],[0,0,0],[0,0,0],[0,0,0]]

a=[10,40,70]
dp[0]=a
print(dp)
#出力 [[10,40,70],[0,0,0],[0,0,0],[0,0,0],[0,0,0]]


# 以下が問題のコード#######
dp[0][1]=60
print(dp)
#出力 [[10,40,70],[0,60,0],[0,60,0],[0,60,0],[0,60,0]]

本当はdp[0][1]の所だけに60を代入したかった．(以下)

print(dp)
#出力 [[10,40,70],[0,60,0],[0,0,0],[0,0,0],[0,0,0]]

原因

以下のようにリストを初期化してしまうと，要素であるリストが全て同じオブジェクトとして生成されてしまう．(以下)

# ダメな例1
dp=[[0,0,0]]*5

# ダメな例2
dp=[[0]*3]*5

解決法

以下のように，「内包表記」 でリストを初期化すれば，全て異なるオブジェクトとして生成できる．

dp=[[0]*3 for i in range(5)]
dp[0][1]=60
print(dp)
#出力 [[10,40,70],[0,60,0],[0,0,0],[0,0,0],[0,0,0]]

numpyなら楽にできる

import numpy as np

dp=np.zeros((5,3))
dp[0,1]=60
print(dp)
#出力 array([[ 0.,  0.,  0.],[ 0., 60.,  0.],[ 0., 0.,  0.],[ 0., 0.,  0.],[ 0., 0.,  0.]])

まとめ

list型で多次元配列を初期化するなら，内包表記を使う．
ただし，numpyで初期化する方が楽．

参考