Challenge: Find Median in Linear Time

Question:

A straightforward way to find a median from a list of numbers would be to first sort the list and then pick the middle element. Using this approach, the time complexity would be dominated by the sorting algorithm, which is O(n log n). Now, your task is to find a median in a linear time.

Answer:

The main idea of the code below is to use the divide and conquer technique that is also used in Quicksort. The function kth_smallest below resembles many similarities with Quicksort.

import random

# First create a function to get the k-th smallest
# element in an array in linear time
def kth_smallest(arr, k):
  # k here is a 0-based index
  n = len(arr)
  if n == 1 and k==0:
    return arr[0]

  pivot = arr[random.randint(0, n-1)]

  lows, pivots, highs = [], [], []
  for i in range(n):
    if arr[i] < pivot:
      lows.append(arr[i])
    elif arr[i] > pivot:
      highs.append(arr[i])
    else:
      pivots.append(arr[i])
  
  # note: k here is a 0-based index
  if k < len(lows):
    return kth_smallest(lows, k)
  elif k < len(lows) + len(pivots):
    return pivot
  else:
    return kth_smallest(highs, k-len(lows)-len(pivots))

def median_linear_time(arr):
  n = len(arr)
  if n <= 0:
    return None
  
  if n % 2 == 1:
    return kth_smallest(arr, n//2)
  else:
    return (kth_smallest(arr, n//2) + kth_smallest(arr, n//2-1))/2
  return 0

# something to test
input = [0,1,2] # output = 1
print(f'{input}: {median_linear_time(input)}')

input = [0,1,2,3] # output = 1.5
print(f'{input}: {median_linear_time(input)}')

input = [3,1,2,0] # output = 1.5
print(f'{input}: {median_linear_time(input)}')

# compare our implementation with statistics'
input = [6,8,5,4,0,6,5,2,7,0,5,1,0,5,8,2,4,7,4]
import statistics
print(f'{input}: {median_linear_time(input)} vs {statistics.median(input)}')

👍 Have fun while coding with Python!

👌 Please subscribe to receive notifications on future challenges.

💬 If you have any questions simply write a comment down below.

Leave a Comment