如何将numpy中的这两层嵌套for循环改为向量化操作？

实现方法 #

完全可以通过numpy的向量化索引操作实现，不需要循环，以下是两种常用的实现方案：

方案1：使用np.put_along_axis（最简洁） #

这个API是numpy专门为这类沿指定轴匹配索引赋值的场景设计的，代码可读性最高：

import numpy as np

my_array = np.zeros((42, 123, 2021))
assert another_array.shape == (42, 123)

# 替代两层for循环的向量化实现
np.put_along_axis(my_array, another_array[..., np.newaxis], values=1, axis=-1)

这里给another_array加最后一维是为了匹配my_array的维度数，符合put_along_axis的入参要求。

方案2：使用高级索引手动构造索引 #

如果需要更灵活的自定义索引逻辑，可以手动构造前两维的索引数组：

i, j = np.indices(another_array.shape)
my_array[i, j, another_array] = 1

np.indices会直接生成和another_array形状一致的前两个维度的网格索引，三个索引数组逐个位置对应原循环里的i、j、another_array[i,j]，批量完成赋值。

两种方案的效果和你原有的双层循环完全一致，运行效率远高于Python原生循环，数组维度越大性能优势越明显。

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