如何访问Numpy三维数组的短对角线元素并赋值？

嘿，我明白你遇到的困扰了——你想选中三维数组里所有满足i==j、i==k或者j==k的元素，但之前的代码只抓到了i==j==k的体对角线元素对吧？别担心，咱们用正确的掩码逻辑就能搞定这个问题。

步骤1：生成三维索引网格 #

首先，我们需要获取数组每个位置的三个维度索引i、j、k，用np.indices()可以轻松生成对应形状的索引网格：

import numpy as np
N = 5  # 假设你的数组是N×N×N的立方体，这里用5做示例
arr = np.random.rand(N, N, N)  # 创建一个测试用的随机数组

# 生成三个维度的索引矩阵
i, j, k = np.indices(arr.shape)

步骤2：创建正确的掩码 #

你之前的问题应该是误用了逻辑与&，而我们需要的是**逻辑或|**来覆盖三个条件中的任意一种：

# 掩码：满足i==j 或 i==k 或 j==k的所有位置
mask = (i == j) | (i == k) | (j == k)

这个掩码会选中三个对角平面上的所有元素：

• i==j：x-y平面上的对角线（沿z轴延伸）
• i==k：x-z平面上的对角线（沿y轴延伸）
• j==k：y-z平面上的对角线（沿x轴延伸）

步骤3：设置特定值 #

有了掩码之后，直接用它来索引数组并赋值就行：

# 把选中的元素设置为你想要的值，比如0
arr[mask] = 0

验证一下 #

如果你想确认是否选中了正确的元素，可以打印这些元素看看：

print("被选中的短对角线元素：", arr[mask])

扩展：非立方体数组也适用 #

如果你的数组不是立方体（比如形状是(M, N, P)），这个方法同样有效，np.indices()会自动生成对应形状的索引网格，掩码逻辑不需要修改。

内容的提问来源于stack exchange，提问作者user1443613