如何填充numpy數組的「部分對角軸」

Question

我有一個三維的numpy數組A。我想要

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