Python实现两个排序数组和的第k小元素

问题描述

给定两个排序数组A，B,定义集合sum=a+b，a来自A,b来自B,求sum中第k小元素
问题示例

问题示例

输入A=[1,7,11] B=[2,4,6]
输出：k=3 返回7，k=4 返回9
此时，sum={3，5，7，9，11，13，15，17}

代码实现

import heapq

def kth_smallest_sum(A, B, k):
    if not A or not B:
        return None
    # 保存所有访问记录，(和,x,y),初始元素为A和B中第一个元素的和
    minheap = [(A[0] + B[0], 0, 0)]
    visited = set()
    num = None

    for _ in range(k):
        # 利用堆的操作，每次从列表中弹出最小的值和对对应的下标
        num, x, y = heapq.heappop(minheap)
        # 如果数组A有下一个元素，并且A[x+1]和B[y]的组合值未计算过
        if x + 1 < len(A) and (x + 1, y) not in visited:
            # #将计算结果压入堆中
            heapq.heappush(minheap, (A[x + 1] + B[y], x + 1, y))
            # 记录访问记录
            visited.add((x + 1, y))
        # 同理处理A[x]和B[y]
        if y + 1 < len(B) and (x, y + 1) not in visited:
            heapq.heappush(minheap, (A[x] + B[y + 1], x, y + 1))
            visited.add((x, y + 1))

    return num


if __name__ == '__main__':
    generator_A = [1, 7, 11]
    generator_B = [2, 4, 6]
    k = 4
    print("输入：", generator_A, generator_B)
    print("k:", k)
    print("输出：", kth_smallest_sum(generator_A, generator_B, k))