如何在以索引为条件的numpy中创建掩码？

问题描述

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的方式来实现此目标

解决方法

您可以这样做：

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

,

我建议采用以下方式：

tiled = np.tile(np.identity(N),(N,1))
for i in range(N):
    tiled[i*N+i,i] = 0
mask = np.reshape(tiled,N,N))

在这里，您首先创建2d数组，垂直插入零，然后将其重塑为3d数组。代码执行速度比原始代码快。

原文：13.5 µs ± 305 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)

新功能：10.2 µs ± 396 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)