Problem statement

https://binarysearch.com/problems/Merging-K-Sorted-Lists/

Solution

Variation of Leetcode 0023. Merge k Sorted Lists, but for usual lists, not linked ones.

Complexity

It is O(n log k) for time and O(n) for space, where n is the total length of all lists.

Code

class Solution:
    def solve(self, lists):
        heap, ans = [], []

        for i in range(len(lists)):
            if lists[i]: heappush(heap, (lists[i][0], i, 0))

        while heap:
            a, i, j = heappop(heap)
            ans += [a]
            if j + 1 < len(lists[i]):
                heappush(heap, (lists[i][j + 1], i, j + 1))
        
        return ans