python - 访问 Numpy 3 维数组中的短对角线元素

Question

我有一个 3 维 numpy 数组，我想访问它的短对角线元素。假设 i、j、k 是三个维度。是否可以访问 i==j or i==k or j= 的元素=k，这样我就可以将它们设置为特定值。

我尝试通过创建索引掩码变量来解决这个问题。该索引掩码变量被馈送到最终数组，其中 {i=j 或 i=k 或 j=k} 的值设置为特定值。不幸的是，这段代码返回的是 {i=j=k}

的集合

import numpy as np

N = 3

maskXY = np.eye(N).reshape(N,N,1)

maskYZ = np.eye(N).reshape(1,N,N)

maskXZ = np.eye(N).reshape(N,1,N)

maskIndices = maskXY * maskYZ*maskXZ

#set the values of final array using above mask

finalArray[maskIndices] = #specific values

Answer 1

方法#1

我们可以使用 np.ix_ 创建开放网格使用覆盖输入数组维度的范围数组，然后在与问题中描述的语法非常接近的数组之间执行 OR-ing，如下所示 -

i,j,k = np.ix_(*[np.arange(r) for r in finalArray.shape])
mask = (i==j) | (i==k) | (j==k)
finalArray[mask] = # desired values

方法#2

看来，我们还可以按照问题中发布的代码并使用 bool 版本的掩码，然后执行 OR 运算以获得等效的 mask ，就像这样 -

mask = (maskXY==1) | (maskYZ==1) | (maskXZ==1)

但是，这涉及 2D(压缩时)的掩模，因此不会像之前处理 1D 数组的方法那样节省内存。