求助：将数组a垂直相邻元素对映射为数组b对应值的向量化实现

问题背景 #

给定两个NumPy数组：

• 数组a：各列元素为1、2、3的随机值，示例如下：

import numpy as np
a = np.array([[2, 3, 1, 3],
              [3, 2, 1, 3],
              [1, 1, 1, 2],
              [1, 3, 2, 3],
              [3, 3, 1, 3],
              [2, 1, 3, 2]])

• 数组b：前两列包含1、2、3的全部9种有序元素对，第三列为对应元素对的关联非负整数，示例如下：

b = np.array([[1, 1, 6],
              [1, 2, 0],
              [1, 3, 9],
              [2, 1, 6],
              [2, 2, 0],
              [2, 3, 4],
              [3, 1, 1],
              [3, 2, 0],
              [3, 3, 8]])

需要生成数组c，将a中每列的垂直相邻元素对替换为b中匹配元素对的第三列值。比如a第一列的元素序列2、3、1、1、3、2，对应的相邻对是(2,3)、(3,1)、(1,1)、(1,3)、(3,2)，对应b中的值就是4、1、6、9、0，也就是c的第一列。

预期输出c：

c = np.array([[4, 0, 6, 8],
              [1, 6, 6, 0],
              [6, 9, 0, 4],
              [9, 8, 6, 8],
              [0, 1, 9, 0]])

要求实现高效的向量化解决方案，避免循环以提升大量执行时的性能。

向量化解决方案 #

核心思路是将有序元素对转换为唯一索引，通过索引直接映射获取对应值，完全规避Python层面循环。

步骤1：构建元素对映射表 #

由于元素取值为1-3，可将有序对(x,y)转换为(x-1)*3 + (y-1)的索引（减1后取值范围为0-2，索引范围0-8，正好对应b的9行）。基于此构建二维映射数组：

# 构建3x3的映射表，直接通过(x-1,y-1)索引获取对应值
mapping = np.zeros((3, 3), dtype=int)
mapping[b[:,0]-1, b[:,1]-1] = b[:,2]

步骤2：提取相邻元素对并映射取值 #

提取a的上半部分（除最后一行）和下半部分（除第一行）作为元素对的前后元素，直接通过映射数组批量取值：

# 转换为0起始的索引值
x = a[:-1] - 1
y = a[1:] - 1
# 向量化索引取值生成结果数组c
c = mapping[x, y]

完整代码 #

import numpy as np

# 输入数组
a = np.array([[2, 3, 1, 3],
              [3, 2, 1, 3],
              [1, 1, 1, 2],
              [1, 3, 2, 3],
              [3, 3, 1, 3],
              [2, 1, 3, 2]])

b = np.array([[1, 1, 6],
              [1, 2, 0],
              [1, 3, 9],
              [2, 1, 6],
              [2, 2, 0],
              [2, 3, 4],
              [3, 1, 1],
              [3, 2, 0],
              [3, 3, 8]])

# 构建映射表
mapping = np.zeros((3, 3), dtype=int)
mapping[b[:,0]-1, b[:,1]-1] = b[:,2]

# 生成结果数组c
x = a[:-1] - 1
y = a[1:] - 1
c = mapping[x, y]

# 输出验证
print(c)

结果验证 #

运行代码后输出与预期完全一致：

[[4 0 6 8]
 [1 6 6 0]
 [6 9 0 4]
 [9 8 6 8]
 [0 1 9 0]]

性能说明 #

该方案完全基于NumPy底层C实现的向量化操作，无Python循环，在大规模数组重复执行场景下，性能远高于循环实现。

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