如何基于NumPy高效实现双数组等量选元素并最大化唯一值？

问题背景 #

现有两个长度为N的数组，元素取值范围为[0, N-1]，示例如下：

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

需要生成新数组c，满足以下要求：

• c恰好包含来自a和b的各N/2个元素（忽略N为奇数的情况）；
• c中的唯一数值尽可能多，重复值尽可能少；
• 元素位置与原数组一致：即c[i]要么等于a[i]，要么等于b[i]。

已有递归遍历的直接解法：

def traverse(i, a, b, path, n_a, n_b, best, best_path):
    if n_a == 0 and n_b == 0:
        score = len(set(path))
        return (score, path.copy()) if score > best else (best, best_path)

    if n_a > 0:
        path.append(a[i])
        best, best_path = traverse(i + 1, a, b, path, n_a - 1, n_b, best, best_path)
        path.pop()
    
    if n_b > 0:
        path.append(b[i])
        best, best_path = traverse(i + 1, a, b, path, n_a, n_b - 1, best, best_path)
        path.pop()

    return best, best_path

调用示例：

>>> score, best_path = traverse(0, a, b, [], 2, 2, 0, None)
>>> score, best_path
(4, [2, 1, 3, 0])

核心疑问 #

是否可以通过NumPy以更向量化、更高效的方式实现该需求？

解决方案：NumPy向量化优化思路 #

可以通过NumPy实现更高效的解法，核心是贪心策略+向量化计算，彻底规避递归遍历的指数级复杂度（递归时间复杂度为C(N, N/2)，N稍大就完全不可用）。

具体步骤 #

1. 优先选新值：对每个位置，优先选择能引入未出现过的数值的选项，最大化唯一值数量；
2. 严格控数量：选择过程中时刻保证已选a/b的数量不超过N/2；
3. 批量补缺口：第一轮选完后，再批量填充剩余的选择名额，满足数量要求。

实现代码 #

import numpy as np

def greedy_max_unique(a, b):
    N = len(a)
    k = N // 2
    # 初始化：标记选a(1)还是b(0)，已选计数，已出现的值集合
    selected = np.zeros(N, dtype=int)
    count_a = 0
    count_b = 0
    seen = set()
    
    # 第一轮：优先选能带来新值的选项
    for i in range(N):
        val_a, val_b = a[i], b[i]
        can_take_a = count_a < k
        can_take_b = count_b < k
        
        a_is_new = val_a not in seen
        b_is_new = val_b not in seen
        
        if can_take_a and can_take_b:
            if a_is_new and not b_is_new:
                selected[i] = 1
                count_a += 1
                seen.add(val_a)
            elif b_is_new and not a_is_new:
                selected[i] = 0
                count_b += 1
                seen.add(val_b)
            elif a_is_new and b_is_new:
                # 两个都是新值，优先选a（可根据需求调整）
                selected[i] = 1
                count_a += 1
                seen.add(val_a)
            else:
                # 都不是新值，留到第二轮处理
                continue
        elif can_take_a and a_is_new:
            selected[i] = 1
            count_a += 1
            seen.add(val_a)
        elif can_take_b and b_is_new:
            selected[i] = 0
            count_b += 1
            seen.add(val_b)
        else:
            continue
    
    # 第二轮：填充剩余名额
    remaining_a = k - count_a
    remaining_b = k - count_b
    
    if remaining_a > 0:
        # 找到未选择且能选a的位置
        mask = (selected == 0) & (count_a < k)
        idx = np.where(mask)[0][:remaining_a]
        selected[idx] = 1
        count_a += len(idx)
    
    if remaining_b > 0:
        mask = (selected == 1) & (count_b < k)
        idx = np.where(mask)[0][:remaining_b]
        selected[idx] = 0
        count_b += len(idx)
    
    # 生成最终数组
    c = np.where(selected == 1, a, b)
    return len(np.unique(c)), c

# 测试示例
a = np.array([0, 1, 2, 0])
b = np.array([2, 0, 3, 3])
score, c = greedy_max_unique(a, b)
print(f"得分：{score}，结果数组：{c}")

补充说明 #

• 该贪心解法时间复杂度为O(N)，远快于递归，适合处理大尺寸数组；
• 贪心策略不一定能保证全局最优，但绝大多数场景下能得到接近最优的结果；
• 如果必须确保全局最优，可以结合动态规划+NumPy向量化优化，实现复杂度会高于贪心，但仍比递归高效。

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