Skill Set - K-Way Merge

Skills for tackling LeetCode problems - K-Way Merge.

We can use Min Heap to make a sorted traversal of all elements in multiple arrays.

1. Merge K Sorted Lists

Given an array of K sorted LinkedLists, merge them into one sorted list.

# Time complexity O(N log K). Space complexity O(K).
from heapq import *

# class ListNode:
#   def __init__(self, value):
#     self.value = value
#     self.next = None

#   # used for the min-heap
#   def __lt__(self, other):
#     return self.value < other.value

def merge_lists(lists):
    minHeap = []
    
    for root in lists: # add the first element of each lists into the heap
        if root:
            heappush(minHeap, root)
    
    newHead, newTail = None, None
    while minHeap:
        node = heappop(minHeap)
        if not newHead:
            newHead = newTail = node
        else:
            newTail.next = node
            newTail = newTail.next
            
        if node.next:
            heappush(minHeap, node.next)
    return newHead
            

2. Kth Smallest Number in M Sorted Lists

Given M sorted arrays, find the Kth smallest number among all the arrays.

# Time complexity O(K log M). Space complexity O(M). M is the total number of input arrays.

def find_Kth_smallest(lists, k):
    minHeap = []
    for i in range(len(lists)): # add the first element of each list
        heappush(minHeap, (lists[i][0], 0, lists[i]))
    Count, Num = 0,0
    
    while minHeap:
        Num, index, List = heappop(minHeap)
        Count += 1
        if Count == k: # find k smallest numbers
            break
        if len(List) > index + 1:
            heappush(minHeap, (List[index+1], index+1, List))
            
    return Num

3. Kth Smallest Number in a Sorted Matrix

Given an N by N matrix where each row and column is sorted in ascending order, find the Kth smallest element in the matrix.

# Time complexity O(min(K, N) + K log N). Space complexity O(N).
def find_Kth_smallest(matrix, k):
    minHeap = []
    
    for i in range(min(k, len(matrix))):
        heappush(minHeap, (matrix[i][0],0, matrix[i]))
        
    Count, Num = 0,0 
    while minHeap:
        Num, index, row = heappop(minHeap)
        Count += 1
        if Count == k:
            break
        if len(row) > index + 1:
            heappush(minHeap, (row[index+1], index+1, row))
    return Num

Alternative approach: Binary Search.

In each round we use binary search to find the middle number between the range start = matrix[0][0] and end = matrix[n-1][n-1] (not necessary a number in the matrix). With this middle number, we count the numbers smaller than the middle and the numbers greater than the middle. At the same time, we keep track of the biggest number less than or equal to the middle (n2) and the smallest number greater than the middle (n1). If the count if equal to K, n1 will be the number. If the count is less than K, we have to search in the higher part of the matrix: update start = n2. If the count is greater than K, we search in the lower part of the matrix: update end = n1.

# Time complexity O(N log (max - min)). Binary search takes O(log (max - min)). Counting takes O(N).
# Space complexity O(1).
def find_Kth_smallest(matrix, k):
    n = len(matrix) # number of lists
    l, r = matrix[0][0], matrix[n-1][n-1]
    
    while l < r: # binary search
        m = l + (r - l) // 2
        n1, n2 = (matrix[0][0], matrix[n-1][n-1])
        
        count, n1, n2 = counting(matrix, m, n1, n2) # count numbers <= m and > m
        
        if count == k: # exactly k numbers <= m
            return n1
        if count < k: # need to search in a higher part
            l = n2
        else: # need to search in a lower part
            r = n1
            
    return l
            
def counting(matrix, m, n1, n2): 
    count, n = 0, len(matrix)
    row, col = n - 1, 0
    while row >= 0 and col < n:
        if matrix[row][col] > m: # keep track of the smallest number greater than the middle
            n2 = min(n2, matrix[row][col])
            row -= 1
        else: # keep track of the greatest number smaller than the middle
            n1 = max(n1, matrix[row][col])
            count += row + 1
            col += 1
    return count, n1, n2

4. Smallest Number Range

Given M sorted arrays, find the smallest range that includes at least one number from each of the M lists.

# Time complexity O(N log M) where N is the total number of elements in all M arrays.
# Space complexity O(M).
from heapq import *
import math

def find_smallest_range(lists):
    minHeap = []
    start, end = 0, math.inf
    MAX = -math.inf
    
    for arr in lists: # add the first element of each list into the heap
        heappush(minHeap, (arr[0], 0, arr))
        MAX = max(MAX, arr[0]) # keep track of the max value
        
    while len(minHeap) == len(lists):
        num, index, arr = heappop(minHeap)
        if end - start > MAX - num: # find the smallest range
            start = num
            end = MAX
            
        if len(arr) > index + 1: # add more elements to heap
            heappush(minHeap, (arr[index+1], index + 1, arr))
            MAX = max(MAX, arr[index+1])
            
    return [start, end]