如何用NumPy向量化实现这段嵌套循环代码？

原始代码 #

import numpy as np

data = np.array([
    [[5, 4], [1, 0]],
    [[7, 4], [0, 2]],
])

print(data.shape)

req_coordinates = np.array([
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ],
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ],
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ]
])

print(req_coordinates.shape)

data_new = np.empty_like(data)
print(data_new.shape)

for i in range(data.shape[0]):
    for j in range(data.shape[1]):
        for k in range(data.shape[2]):
            #print(i, j, k)
            s = int(req_coordinates[0, i, j, k])
            t = int(req_coordinates[1, i, j, k])
            u = int(req_coordinates[2, i, j, k])
            #print(s, t, u)
            data_new[s, t, u] = data[i, j, k]

print(data)

向量化实现方案 #

你可以利用NumPy的高级索引直接替换嵌套循环，这是NumPy中实现向量化操作的核心高效方式，完全规避Python层面的循环开销。

具体操作：

1. 从req_coordinates中拆分出三个维度的索引数组，它们的形状和data完全匹配，刚好对应data_new的三个维度索引。
2. 通过这组索引数组直接完成批量赋值，无需逐元素循环。

向量化后的代码：

import numpy as np

data = np.array([
    [[5, 4], [1, 0]],
    [[7, 4], [0, 2]],
])

req_coordinates = np.array([
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ],
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ],
    [
        [[0, 1], [0, 1]],
        [[0, 1], [0, 1]]
    ]
])

data_new = np.empty_like(data)
# 拆分三个维度的索引数组
s_idx, t_idx, u_idx = req_coordinates
# 一次性完成所有赋值
data_new[s_idx, t_idx, u_idx] = data

print(data_new)

补充说明 #

• numpy.vectorize并非真正的向量化实现，底层仍会执行循环，不适合这个场景；numpy.put更适配一维数组的批量赋值，多维复杂索引场景下高级索引是更直接的最优解。

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