如何无需for循环实现Numpy数组映射：new_array[i][j] = b[a[i][j]]

解决方法 #

你需要的操作是Numpy原生支持的整数数组索引，直接把数组a作为索引传入b即可得到符合要求的结果，全程无循环、纯向量化实现：

import numpy as np

a = np.array([[0, 2, 0], [0, 2, 0]])
b = np.array([0, 0, 2])
# 直接用a作为索引取值
new_array = b[a]

执行后new_array的输出完全匹配你的预期：

array([[0, 2, 0],
       [0, 2, 0]])

原理解释 #

当你将整数型Numpy数组作为索引传入另一个数组时，Numpy会自动遍历索引数组的每个位置，将该位置的元素值作为下标取出被索引数组的对应值，最终返回的数组形状和索引数组完全一致，天然满足new_array[i][j] = b[a[i][j]]的逻辑，执行效率远高于手写嵌套循环，数组越大性能优势越明显。

你也可以用np.take函数实现完全相同的效果：

new_array = np.take(b, a)

注意事项 #

索引数组a中的所有元素值必须小于被索引数组b的长度，否则会触发索引越界错误，该要求和手写循环的逻辑一致。

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