如何在含重复值的两个Numpy数组间实现快速索引映射？

问题背景 #

尝试使用Numpy的recfunctions.join_by执行数组连接时，发现该方法无法处理重复值。需要基于连接列构建两个数组间的索引映射，要求：

• 完全基于Numpy实现，利用向量化操作
• 保留原数组中的重复值，映射长度需与输入数组x一致
• 不使用原生Python循环、Pandas或numpy-indexed库
• 无匹配项的位置用掩码（如-1）标记

映射关系需满足：

mapping = f(x, y)
x = y[mapping]  # 无匹配项位置可通过掩码过滤

示例需求 #

x = np.array([1,1,2,100])
y = np.array([1,2,3,4,5,6,7])
# 期望映射结果：[0, 0, 1, -1]

现有方案的问题 #

• np.searchsorted无法处理x中的重复值，输出完全不符合需求：

import numpy as np

x = np.array([1,1,2,100])
y = np.array([1,2,3,4,5,6,7])

mapping = np.searchsorted(x, y)
# 输出：[0 2 3 3 3 3 3]

• np.where(np.isin(x,y))会自动去重，得到的映射长度与x不一致：

import numpy as np

x = np.array([1,1,2,100])
y = np.array([1,2,3,4,5,6,7])

mapping = np.where(np.isin(x, y))[0]
# 输出：[0, 1, 2]

纯Numpy解决方案 #

实现代码 #

import numpy as np

def get_mapping(x, y):
    # 获取y的唯一值及每个值在y中的首次出现索引
    y_unique, idx_y = np.unique(y, return_index=True)
    # 查找x元素在y_unique中的位置
    pos = np.searchsorted(y_unique, x)
    # 生成掩码：标记x中不在y里的元素
    mask = (pos >= len(y_unique)) | (y_unique[pos] != x)
    # 构建映射数组，无匹配项设为-1
    mapping = idx_y[pos]
    mapping[mask] = -1
    return mapping

# 测试示例
x = np.array([1,1,2,100])
y = np.array([1,2,3,4,5,6,7])
mapping = get_mapping(x, y)
print(mapping)  # 输出：[0 0 1 -1]

# 验证映射结果
valid_mask = mapping != -1
result = np.full_like(x, np.nan, dtype=np.float64)
result[valid_mask] = y[mapping[valid_mask]]
print(result)  # 输出：[  1.   1.   2. nan]

方案说明 #

1. 唯一值提取：通过np.unique(y, return_index=True)获取y的唯一值数组，以及每个唯一值在y中的首次出现索引，确保重复值的匹配一致性
2. 位置查找：使用np.searchsorted在有序的唯一值数组中查找x元素的位置，避免原数组重复值的干扰
3. 掩码校验：通过掩码确认查找结果的准确性，排除searchsorted返回的临近不匹配位置
4. 映射构建：将有效位置映射到y的原始索引，无匹配项标记为-1，保证映射长度与x完全一致

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