如何对用两个2D数组索引4D NumPy数组的操作进行向量化优化

你可以通过 NumPy 高级索引实现完全向量化的赋值，无需嵌套循环：

实现逻辑 #

目标数组的四个维度为 (T, M, N, D)，我们只需要构造与源数组shape=(T, D)完全匹配的四组索引数组，分别对应四个维度的赋值位置即可：

• 第一维（T轴）索引：通过np.arange(T)[:, None]生成，广播后形状为(T, D)，每个位置对应对应时间步t
• 第四维（D轴）索引：通过np.arange(D)[None, :]生成，广播后形状为(T, D)，每个位置对应对应特征维度d
• 第二、三维索引直接使用已有的index_dim_1和index_dim_2即可，二者本身形状就是(T, D)，无需额外处理

完整代码 #

import numpy as np

T, M, N, D = 5, 4, 3, 2  # 示例维度定义
target = np.zeros((T, M, N, D))  # 目标数组
source = np.arange(T*D).reshape(T, D)  # 示例源数组，可替换为任意符合形状的数值
index_dim_1 = np.array([  # 形状为(T,D)的第二轴索引，取值范围0~M-1
    [0, 3],
    [1, 2],
    [2, 1],
    [3, 0],
    [0, 3]
])
index_dim_2 = np.array([  # 形状为(T,D)的第三轴索引，取值范围0~N-1
    [0, 2],
    [1, 1],
    [2, 0],
    [0, 2],
    [1, 1]
])

# 向量化赋值核心逻辑
idx_t = np.arange(T)[:, None]
idx_d = np.arange(D)[None, :]
target[idx_t, index_dim_1, index_dim_2, idx_d] = source

该实现和原有循环的输出结果完全一致，在大维度场景下运行效率会有数量级的提升。

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