如何按值与索引规则配对两个长度相同的一维NumPy数组

实现思路 #

该需求本质是构建二部图的连通分量：将数组a的唯一值、数组b的唯一值作为二部图的两组顶点，每一组索引对应的a值和b值之间连一条边，最终每个连通分量内的a值集合与b值集合即为期望的配对结果。

完整代码 #

import numpy as np
from collections import defaultdict, deque

# 示例输入
a = np.array([1, 1, 1, 2, 2, 3, 4,  5])
b = np.array([7, 7, 8, 8, 9, 8, 10, 10])

# 构建双向映射关系
a_to_b = defaultdict(set)
b_to_a = defaultdict(set)
for a_val, b_val in zip(a, b):
    a_to_b[a_val].add(b_val)
    b_to_a[b_val].add(a_val)

visited_a = set()
visited_b = set()
result = []

# BFS遍历找连通分量
for a_val in sorted(a_to_b.keys()):
    if a_val not in visited_a:
        q = deque([('a', a_val)])
        current_a = set()
        current_b = set()
        while q:
            type_, val = q.popleft()
            if type_ == 'a' and val not in visited_a:
                visited_a.add(val)
                current_a.add(val)
                for b_val in a_to_b[val]:
                    if b_val not in visited_b:
                        q.append(('b', b_val))
            elif type_ == 'b' and val not in visited_b:
                visited_b.add(val)
                current_b.add(val)
                for a_v in b_to_a[val]:
                    if a_v not in visited_a:
                        q.append(('a', a_v))
        # 排序后存入结果
        result.append((sorted(current_a), sorted(current_b)))

# 输出结果
for a_group, b_group in result:
    print(f"{a_group} 与 {b_group} 配对")

运行结果 #

[1, 2, 3] 与 [7, 8, 9] 配对
[4, 5] 与 [10] 配对

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