如何在不使用嵌套 for 循环的情况下根据其索引值划分 2d numpy 数组中的每个元素？

Question

我正在尝试提高从 MATLAB 转换为 Python 的一些代码的性能，并且发现使用嵌套 for 循环如下所示：

    arr1 = np.zeros((NN, MM))
    for i in range(NN):
        for j in range(MM):
            arr1[i, j] = arr[i, j]/(1+np.abs(i-j))

导致脚本运行几秒钟。由于我将在大型数据集上使用此代码，这可能会不必要地浪费数小时的计算时间。

有什么方法可以用vectorization/broadcasting达到同样的效果吗？

Answer 1

您可以像这样预先计算索引。

import numpy as np
arr1 = np.zeros((NN, MM))

然后使用 numpy indices 函数，它 returns 一个 3d 矩阵，其中第一个索引表示 i 或 j

grid = np.indices(arr1.shape)

然后划分。

arr1 /= (1+np.abs(grid[0]-grid[1]))

Answer 2

如果arr1是np.zeros，那么整个操作都是多余的，呵呵。

但是，是的，可以加快运算速度 - 如果我们关注分母 - (1+np.abs(i-j))，你可以这样分解它：

J = np.tile(np.arange(MM), (NN, 1))

I = np.tile(np.reshape(np.arange(NN), (NN, 1)), (1, MM))

然后让

x = x / (1 + I - J)

但是请注意，当 I = J + 1 时 1 + I - J 将为零。

How to divide each element in a 2d numpy array in relation to its index value without the use of nested for loops?