如何不使用for循环实现二维numpy数组的元素映射？

你的需求是把二维NumPy数组的元素按指定位置映射重组，原来的for循环在大数组场景下效率不足，完全可以用NumPy的向量化索引来实现，全程不需要循环，速度会快很多。

优化后的代码 #

import numpy as np

states = np.array([["a", "b", "c", "d"], ["A", "B", "C", "D"]])
pos_new = np.array([[0, 1], [0, 2]])
pos_old = np.array([[1, 2], [2, 0]])

# 构造行索引，让每一行对应各自的位置索引
row_indices = np.arange(states.shape[0])[:, None]
# 直接通过向量化索引完成赋值
states[row_indices, pos_new] = states[row_indices, pos_old]

print(states)

代码解释 #

1. 先把pos_new和pos_old转换成NumPy数组，这样才能和原数组的索引机制兼容。
2. row_indices = np.arange(states.shape[0])[:, None]生成形状为(2,1)的数组[[0], [1]]，和pos_new（形状(2,2)）结合时会自动广播匹配维度，保证每一行的位置索引对应到正确的行上。
3. 直接通过states[row_indices, pos_new]选中所有要替换的目标位置，然后用states[row_indices, pos_old]的对应值赋值，一步完成所有行的操作，完全替代原来的for循环。

输出结果 #

运行后会得到和原代码完全一致的输出：

[['b' 'c' 'c' 'd']
 ['C' 'B' 'A' 'D']]

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