【python刷题】多个有序数组

def mergeList(A, B):
    s1 = len(A)
    s2 = len(B)
    i,j = 0,0
    res = []
    while i < s1 and j < s2:
        if A[i] <= B[j]:
            res.append(A[i])
            i += 1
        else:
            res.append(B[j])
            j += 1
    res = res + A[i+1:] + B[j+1:]
    return res

A = [1,2,5,7,9]
B = [2,4,6,8,10,11,34,55]
res = mergeList(A, B)
print(res)

def mergeMultiList(lists):

    import heapq
    from collections import deque
    lists = list(map(lambda x: deque(x), lists))
    pq = []
    for ind, val in enumerate(lists):
        pq.append((val.popleft(), ind))
    heapq.heapify(pq)
    res = []
    while pq:
        value, index = heapq.heappop(pq)
        print(value, index)
        res.append(value)
        if lists[index]:
            heapq.heappush(pq, (lists[index].popleft(), index))
    return res

lists = [[1,2,5,7,9],[2,4,6,8,10,11,34,55],[1,3,5,8,10,15]]
res = mergeMultiList(lists)
print(res)

class Solution:

    """
    @param A: An integer array.
    @param B: An integer array.
    @return: a double whose format is *.5 or *.0
    """

    def findMedianSortedArrays(self, A, B):
        n = len(A) + len(B)
        if n % 2 == 1:
            return self.findKth(A, B, n / 2 + 1)
        else:
            smaller = self.findKth(A, B, n / 2)
            bigger = self.findKth(A, B, n / 2 + 1)
            return (smaller + bigger) / 2.0

    def findKth(self, A, B, k):
        if len(A) == 0:
            return B[int(k - 1)]
        if len(B) == 0:
            return A[int(k - 1)]
        if k == 1:
            return min(A[0], B[0])

        a = A[int(k / 2) - 1] if len(A) >= k / 2 else None
        b = B[int(k / 2) - 1] if len(B) >= k / 2 else None

        if b is None or (a is not None and a < b):
            return self.findKth(A[int(k / 2):], B, int(k - k // 2))
        return self.findKth(A, B[int(k / 2):], int(k - k // 2))

s = Solution()
print(s.findMedianSortedArrays([1, 2, 3, 4, 5, 6], [2, 3, 4, 5]))