educative.io

Kth Smallest Number in M Sorted Lists reduce space complexity

Hi Team,
Instead of storing each list again in min-heap, if store the length of the array and moving index of the array. it reduces space complexity to O(k). please comment on this below.

new min heap parameters array_index,length, instead of list.

def find_Kth_smallest(lists,k):
    min_heap=[]
    for i in range(len(lists)):
        heappush(min_heap,(lists[i][0],0,i,len(lists[i])))
    counter=0
    while len(min_heap):
        #print(min_heap)
        num,i,array_index,length=heappop(min_heap)
        counter+=1
        if counter==k:
            break;
        if length>i+1:
            heappush(min_heap,(lists[array_index][i+1],i+1,array_index,length))
    return num

Your solution doesn’t change the space complexity. As long as you have M elements to track in the heap the space complexity is O(M).

The tuples in the original solution contain references to the input lists, no list creation is needed. It takes O(1) space for each list reference which is the same as tracking list index or length (you are tracking both but you only need one of them).

1 Like

you can reduce space complexity by not using a heap:

def find_Kth_smallest(lists, k):
  number = -1
  while lists:
    current_min = float('inf')
    current_idx = None
    for idx, l in enumerate(lists):
      if l[0] < current_min:
        current_idx = idx
        current_min = l[0]
    k-=1
    if k ==0:
      number = current_min
      break
    
    lists[current_idx].pop(0)
    if not lists[current_idx]:
      lists.pop(current_idx)
    
  return number

though time complexity would be significantly higher because pop(0) ia an O(n) operation.
To solve this you should cast all the lists as collections.deque where you can popleft in O(1)