How to merge k sorted arrays, solution didn't work allowing duplicates.!

Question

Given k sorted arrays of size n each, merge them and print the sorted output.

The algorithm I followed is

• iterate of over each array
  • pick the ith index in k arrays
  • insert() in minheap
  • delMin() and append result array.

from heapq import heappop, heappush

def merge_k_arrays(list_of_lists):
    result = [] #len(list_of_lists[0])*len(list_of_lists)
    minHeap= []
    n, k=0,0

    print(list_of_lists)
    while n < len(list_of_lists[0]):
        if n ==0:# initial k size heap ready
            while k < len(list_of_lists):
                element= list_of_lists[k][n]
                heappush(minHeap ,element )
                k+=1
            result.append(heappop(minHeap))
        else: # one at a time.
            k =0
            while k < len(list_of_lists):
                element = list_of_lists[k][n]
                heappush(minHeap , element)
                result.append(heappop(minHeap))
                k+=1
        n += 1

    # add the left overs in the heap
    while minHeap:
        result.append(heappop(minHeap))

    return result

Input:

input = [   [1, 3, 5, 7],
            [2, 4, 6, 8],
            [0, 9, 10, 11],

        ] 

Output:

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]

input:

input = [   [1, 3, 5, 7],
            [2, 4, 6, 8],
            [3, 3, 3, 3],
            [7, 7, 7,7]
        ]

output:

[0, 1, 2, 3, 3, 3, 4, 5, 6, 3, 7, 7, 7, 7, 3, 7, 8, 9, 10, 11]

Could anyone help me know what piece is missing from my algorithm in order to merge the duplicate arrays in the second input too?

Answer 1

from heapq import *

def mergeSort (arr):
    n = len(arr)
    minHeap = []
    
    for i in range(n):
        heappush(minHeap, (arr[i][0], i, arr[i]))
    
    print(minHeap)
    
    result = []
    while len(minHeap) > 0:
        num, ind, arr = heappop(minHeap)
        result.append(num)
        
        if len(arr) > ind + 1:
            heappush(minHeap, (arr[ind+1], ind+1, arr))
        
    
    return result
    
    

input = [   [1, 3, 5, 7],
            [2, 4, 6, 8],
            [0, 9, 10, 11],
            [-100]

        ] 
        
print(mergeSort(input))

Answer 2

To fix your code, move the result.append(heappop(minHeap)) in your second nested while loop to the outside of the nested while loop, like in your first nested while loop. This will make your code work.

        else: # one at a time.
        k =0
        while k < len(list_of_lists):
            element = list_of_lists[k][n]
            heappush(minHeap , element)

            k+=1
        result.append(heappop(minHeap))
    n += 1

If you have any space constraints, this is still problematic since you are adding nearly your entire input into the heap. If space is not an issue, there is a much cleaner way to write your solution:

def merge(A):
    result = []
    heap = [e for row in A for e in row]
    heapify(heap)
    for i in range(len(heap)):
        result.append(heappop(heap))
    return result

Otherwise, you will need to use a smarter solution that only allows the heap to have k elements in total, with one element from each list, and the new element you add each step should come from the origin list of the element that was just popped.