每日一题 【每日一题】两个排序数组合的第 k 个小元素 -Python-20211018

Jack · 2021年10月18日 · 最后由 Jack 回复于 2021年10月19日 · 14 次阅读
本帖已被设为精华帖!
  1. 问题描述: 给定两个排好序的数组 A,B,定义集合 sum = a + b, 其中 a 来自数组 A,b 来自数组 B,求 sum 中第 k 小的元素
  2. 问题示例: A = [1, 7, 11] , B = [2, 4, 6] sum = [3, 5, 7, 9, 11, 13, 13, 15, 17] 当 k = 3,返回 7 当 k = 4, 返回 9
Jack 将本帖设为了精华贴 10月18日 21:03

参考代码:

class Solution:
    def kthSmallestSum(self, A, B, k):
        if not A or not B:
            return None
        n, m = len(A), len(B)
        minheap = [A[0] + B[0], 0, 0]
        visited = set([0])
        num = None
        for _ in range(k):
            num, x, y = heapq.heappop(minheap)
            if x+1 < n and (x + 1) * m + y not in visited:
                heapq.heappush(minheap, (A[x + 1] + B[y], x + 1, y))
                visited.add((x + 1) * m + y)
            if y+1 < m and x * m + y + 1 not in visited:
                heapq.heappush(minheap, (A[x] + B[y + 1], x, y+1))
                visited.add(x * m + y + 1)
        return num
需要 登录 后方可回复, 如果你还没有账号请点击这里 注册