NumPy数组合并：保留小数组全量值并从大数组补充指定数量唯一值

核心需求说明 #

现有两个NumPy数组，大数组规模可达数十万条，小数组仅数千条且是大数组的子集，合并要求：

• 完整保留小数组全部元素
• 补充元素需为大数组中未在小数组出现过的唯一值
• 最终合并数组长度符合指定要求

实现思路 #

1. 提取大数组中不属于小数组的元素，做去重处理，可选择保留大数组原有顺序或随机抽取
2. 计算需要补充的元素数量，提前做边界校验避免报错
3. 合并小数组与指定数量的补充元素，可根据需要选择合并后的排序规则

完整代码示例 #

import numpy as np

def merge_arr(large_arr: np.ndarray, small_arr: np.ndarray, target_len: int, keep_large_order: bool = True) -> np.ndarray:
    # 提取大数组中不在小数组里的候选补充元素
    supplement_candidates = large_arr[~np.isin(large_arr, small_arr)]
    # 对候选元素去重
    if keep_large_order:
        # 保留大数组中的原有顺序
        _, idx = np.unique(supplement_candidates, return_index=True)
        unique_supplement = supplement_candidates[np.sort(idx)]
    else:
        # 去重后打乱，随机抽取
        unique_supplement = np.random.permutation(np.unique(supplement_candidates))
    
    # 边界校验
    small_len = len(small_arr)
    supplement_need = target_len - small_len
    if supplement_need < 0:
        raise ValueError("目标长度小于小数组长度，无法完整保留小数组全部元素")
    if supplement_need > len(unique_supplement):
        raise ValueError(f"大数组可用补充元素仅{len(unique_supplement)}个，不足补充{supplement_need}个")
    
    # 合并输出，若需要完全匹配大数组原有顺序，替换为下方注释代码
    # return large_arr[np.isin(large_arr, small_arr) | np.isin(large_arr, unique_supplement[:supplement_need])]
    return np.concatenate([small_arr, unique_supplement[:supplement_need]])

# 测试示例
if __name__ == "__main__":
    first_arr = np.array([3, 5, 8, 10, 14, 15, 16, 18])
    second_arr = np.array([5, 10, 15])
    target_len = 5
    # 要和示例输出顺序一致，开启保留大数组原有顺序即可
    merged_arr = merge_arr(first_arr, second_arr, target_len, keep_large_order=True)
    print(merged_arr)

性能说明 #

该方案使用NumPy原生向量化操作实现，np.isin处理100万条规模的大数组仅需毫秒级耗时，完全符合你的数据规模要求。

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