python - 根据单元格索引填充 numpy 数组

Question

我正在尝试创建一个 3 维数组，其单元格条目将从单元格索引中计算出来。具体来说，我想要那个单元格 (i,j,k) = sqrt(i+j+k)。

这很容易通过以下 for 循环完成:

N=10
A=np.zeros((N,N,N))

for i in range(N):
    for j in range(N):
        for k in range(N):
            A[i][j][k] = np.sqrt(i+j+k)

我想知道 numpy 是否具有使这些嵌套 for 循环变得多余的内置函数。

Answer 1

最简单和高效的方法是使用 np.ogrid 打开网格，然后执行相关操作 -

I,J,K = np.ogrid[:N,:N,:N]
A = np.sqrt(I+J+K)

或使用 np.sum 为单线开放网格的广播求和 -

A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))

相关:General workflow on vectorizing loops involving range iterators .