如何在Python中合并3个有序数组并返回元素来源步骤

核心思路 #

采用多指针归并方案，逻辑和归并排序的合并阶段一致，额外增加来源标识记录即可：

• 为每个数组初始化一个指向首个元素的指针
• 每次比较所有指针当前指向的元素，选中值最小的元素
• 将最小元素对应的原数组标识加入结果列表，对应数组的指针后移一位
• 若某个数组的指针超出自身长度，后续不再参与比较
• 循环执行直到所有数组的元素都被处理完成

注：如果多个数组的当前元素值相等，可通过调整数组的预设优先级决定优先选取哪个，优先级规则可自行定义，本实现默认优先级顺序为A > B > C，刚好匹配示例要求。

实现代码 #

# 输入的三个有序数组
A = [1, 3, 4, 6]
B = [2, 3, 4, 5]
C = [1, 5, 9]

# 打包数组和对应标识，调整列表顺序即可修改值相等时的选取优先级
arr_with_tag = [(A, "A"), (B, "B"), (C, "C")]
# 初始化每个数组的指针，初始都指向0位
ptrs = [0] * len(arr_with_tag)
result = []

while True:
    current_min = float("inf")
    selected_arr_idx = -1
    # 遍历所有还有剩余元素的数组，找当前最小值对应的数组
    for i in range(len(arr_with_tag)):
        arr, tag = arr_with_tag[i]
        current_ptr = ptrs[i]
        if current_ptr < len(arr) and arr[current_ptr] < current_min:
            current_min = arr[current_ptr]
            selected_arr_idx = i
    # 所有数组处理完成，退出循环
    if selected_arr_idx == -1:
        break
    # 记录选中的数组标识，对应指针后移
    result.append(arr_with_tag[selected_arr_idx][1])
    ptrs[selected_arr_idx] += 1

print(result)

输出结果 #

运行上述代码得到的输出为：
['A', 'C', 'B', 'A', 'B', 'A', 'B', 'B', 'C', 'A', 'C']
和需求完全匹配。

扩展说明 #

本方案可直接扩展到N个有序数组的合并场景，仅需要在arr_with_tag列表中新增（数组, 标识）元组即可，不需要修改核心逻辑。

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