numpy配列の "部分対角軸"を埋める方法

Question

私はnumpy配列Aの3次元を持っています。私は持っています

A[i, i, :] = 1.0 

すべてiです。

numpyでこれをどのようにエレガントに行うことができますか？すべてのforループを避けることは可能ですか？

Answer 1

アプローチ＃1あなたはnp.eye斜めマスクを作成して割り当てるには、ブール・インデックスを使用することができます -

mask = np.eye(A.shape[0], dtype=bool) 
A[mask] = 1.0 

整数ベースのインデックスを使用したアプローチ＃2 -

r = np.arange(A.shape[0]) 
A[r,r] = 1.0 

アプローチ3＃2Dに変更し、行に割り当てます -

n = A.shape[0] 
A.reshape(-1,A.shape[-1])[np.arange(n)*(n+1)] = 1.0