如何向量化实现基于两个整数向量生成匹配矩阵的操作

问题概述 #

给定两个整数向量A和B，需构建维度为(len(B), len(A))的矩阵X：对于任意索引对(i,j)，当A[i] == B[j]时X[j][i] = 1，否则为0。原循环实现可运行，但在向量规模较大时效率不足，需采用向量化方案优化。

原循环实现代码：

import numpy as np

A = np.array([3,6,12,14])
B = np.array([1,12,14,23,33])
X = np.zeros((len(B), len(A))) 

for i in range(len(A)):
    for j in range(len(B)):
        if A[i] == B[j]:
            X[j][i] = 1

向量化实现方案 #

方案1：基于NumPy广播机制（推荐） #

利用NumPy的广播特性，直接对扩展维度后的数组进行逐元素比较，一步生成目标矩阵：

import numpy as np

A = np.array([3,6,12,14])
B = np.array([1,12,14,23,33])

# 将B转为列向量，与A广播后比较，再转为整数类型
X = (B[:, np.newaxis] == A).astype(int)

核心说明：

• B[:, np.newaxis] 将原形状为(5,)的B转换为(5,1)的列向量
• 与形状为(4,)的A进行==比较时，广播机制会自动将两个数组扩展为(5,4)的维度，逐元素判断是否相等
• astype(int)将布尔值数组（True/False）转换为0/1的整数矩阵

方案2：利用np.isin实现 #

若需更直观的匹配逻辑，可结合np.isin和行重复操作：

# 先判断A中每个元素是否存在于B，再将结果重复len(B)行
X = np.isin(A, B).astype(int)[np.newaxis, :].repeat(len(B), axis=0)

注：此方案效率略低于广播方案，仅作为替代选项。

结果验证 #

两种向量化方案生成的X与原循环代码结果完全一致：

array([[0, 0, 0, 0],
       [0, 0, 1, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])

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