如何填充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