Implement Merge Sort

merge_sort.py

import math

def merge_sort(data, p: int, q: int):
  # If something goes wrong
  if q < p :
    return data
  # Base case
  if q - p == 1:
    if data[q] < data[p]:
      return data[:p] + [data[q], data[p]] + data[q+1:]
    return data
 
  # divide and sort recursively
  sorted_list_1 = merge_sort(data, p, math.floor((p+q)/2))[p: math.floor((p+q)/2 + 1)]
  sorted_list_2 = merge_sort(data, math.floor((p+q)/2)+1, q)[math.floor((p+q)/2)+1: q+1]

  # merge sorted_list_1 and sorted_list_2
  iter_1 = iter(sorted_list_1)
  iter_2 = iter(sorted_list_2)
  result = []
  first = next(iter_1)
  second = next(iter_2)

#   print(p, q, first, second, sorted_list_1, sorted_list_2)
  while True:
    if first < second:
      result.append(first)
      try:
        first = next(iter_1)
      except StopIteration:
        result.append(second)
        result.extend([x for x in iter_2])
        break
    else:
      result.append(second)
      try:
        second = next(iter_2)
      except StopIteration:
        result.append(first)
        result.extend([x for x in iter_1])
        break
  
#   print(data, p, q, result, 'after sorting',data[:p] + result + data[q+1:])
  return data[:p] + result + data[q+1:]

assert merge_sort([5,4,8,7,6,4,9,0], 0, 7) == [0,4,4,5,6,7,8,9]