python - 如何在基于索引的numpy中创建掩码？

Question

N = 5
mask = np.zeros((N, N, N))
for i in range(N):
    for j in range(N):
        for k in range(N):
            if j==k and i!=j:
                mask[i,j,k] = 1

目前我正在按照上面显示的代码进行操作，我觉得必须有一种更有效和pythonic的方式来实现这个目标

标签： pythonnumpymatrixmask

Answer 1

你可以做：

import numpy as np
N = 5
i, j, k = np.ogrid[:N, :N, :N]
mask = (j == k) & (i != j)