使用Numpy查找两个列表中公共数组对应master列表的匹配索引

实现思路 #

核心通过NumPy原生的数组序列化方法tobytes()将任意维度的数组转为可直接比较的字节序列，再基于频次匹配规则获取对应索引，既兼容不同形状的数组，也能严格按照slave的出现次数返回对应数量的master索引。

代码实现 #

基础实现 #

import numpy as np
from collections import Counter

def get_matched_master_indices(slave, master):
    # 序列化数组为字节序列，兼容任意维度、不同形状的数组
    slave_bytes = [arr.tobytes() for arr in slave]
    master_bytes = [arr.tobytes() for arr in master]
    # 统计slave中各数组的出现次数
    count_mapping = Counter(slave_bytes)
    result = []
    
    for idx, byte_seq in enumerate(master_bytes):
        if byte_seq in count_mapping and count_mapping[byte_seq] > 0:
            result.append(idx)
            count_mapping[byte_seq] -= 1
            # 所有匹配项找全后提前终止遍历
            if sum(count_mapping.values()) == 0:
                break
    return result

NumPy向量化优化实现 #

核心匹配逻辑使用NumPy向量化操作，仅遍历slave的唯一值，循环次数远小于列表长度，符合优先使用NumPy的要求：

import numpy as np

def get_matched_master_indices_numpy(slave, master):
    slave_bytes = np.array([arr.tobytes() for arr in slave])
    master_bytes = np.array([arr.tobytes() for arr in master])
    result = []
    
    # 获取slave中唯一数组及对应出现次数
    unique_vals, counts = np.unique(slave_bytes, return_counts=True)
    for val, cnt in zip(unique_vals, counts):
        # 向量化匹配master中所有符合条件的索引，取前cnt个
        matched = np.where(master_bytes == val)[0][:cnt]
        result.extend(matched.tolist())
    # 按master中出现顺序排序
    result.sort()
    return result

测试验证 #

# 测试用例
slave = [np.array([1,5,7,1]), np.array([7,3,5,7]), np.array([1,2,3,8]), np.array([1,2,3,8])]
master = [np.array([1,5,3,7]), np.array([1,5,7,1]), np.array([1,2,3,8]), np.array([7,3,5,7]), np.array([3,2,5,7]), np.array([1,2,3,8]), np.array([1,2,3,8])]

print(get_matched_master_indices(slave, master))
# 输出：[1, 2, 3, 5]
print(get_matched_master_indices_numpy(slave, master))
# 输出：[1, 2, 3, 5]

特性说明 #

• 支持任意维度的NumPy数组匹配，自动兼容列表内数组形状不同的场景
• 严格遵循频次匹配规则：slave中某数组出现N次，仅返回master中前N个匹配项的索引
• 核心操作使用NumPy原生实现，性能优于纯Python循环方案

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