如何移除NumPy多维索引循环？优化数组赋值效率

当然可以用NumPy的高级索引完全替代嵌套循环，实现高效的向量化操作，代码更简洁且运行速度更快。

优化后的代码 #

import numpy as np

x = 3
y = 2
h = 10
w = 12

arr = np.zeros((x, h, w))
index_arr1 = np.array(
    [[[9, 0, 7, 1], [1, 3, 6, 8]],
     [[8, 1, 2, 8], [0, 1, 0, 9]],
     [[0, 8, 7, 8], [0, 4, 5, 9]]]
)
index_arr2 = np.array([4, 8])
value_arr = np.ones((x, y))

# 构造各维度的索引数组
idx_x = np.arange(x)[:, np.newaxis, np.newaxis]  # 形状(x,1,1)，广播适配到(x,y,4)
idx_h = index_arr1  # 直接使用原索引数组，形状(x,y,4)
idx_w = index_arr2[np.newaxis, :, np.newaxis]    # 形状(1,y,1)，广播适配到(x,y,4)

# 向量化赋值
arr[idx_x, idx_h, idx_w] = value_arr[:, :, np.newaxis]

# 验证结果与原循环一致（可选）
arr_loop = np.zeros((x, h, w))
for i in range(y):
    for j in range(x):
        arr_loop[j, index_arr1[j,i], index_arr2[i]] = value_arr[j,i]
print(np.array_equal(arr, arr_loop))  # 输出True

关键逻辑说明 #

1. 索引数组适配：
   • idx_x：把0到x-1的一维数组扩展为三维形状(x,1,1)，通过广播自动匹配index_arr1的(x,y,4)结构，对应原循环中的j维度。
   • idx_w：把index_arr2扩展为(1,y,1)的三维数组，广播后每个index_arr2[i]会和index_arr1[j,i]里的4个索引一一配对，对应原循环中的index_arr2[i]。
2. 值数组适配：
   • value_arr[:, :, np.newaxis]把(x,y)的数组扩展为(x,y,1)，广播后每个value_arr[j,i]会自动填充到index_arr1[j,i]指向的4个位置，和原循环逻辑完全对齐。

效率优势 #

NumPy的向量化操作基于底层C语言实现，避开了Python循环的解释器开销，当x、y或单组索引数量（比如示例中的4）较大时，性能提升会非常显著。

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