Skip to content

Instantly share code, notes, and snippets.

@patricktrainer
Created June 15, 2023 04:01
Show Gist options
  • Save patricktrainer/7309cd2ecc08f7aa2421ada71519223c to your computer and use it in GitHub Desktop.
Save patricktrainer/7309cd2ecc08f7aa2421ada71519223c to your computer and use it in GitHub Desktop.
Calculate the median given a stream of integers from stdin
import heapq
class MedianCalculator:
def __init__(self):
self.heaps = [], []
def add_num(self, num):
small, large = self.heaps
heapq.heappush(small, -heapq.heappushpop(large, num))
if len(large) < len(small):
heapq.heappush(large, -heapq.heappop(small))
def find_median(self):
small, large = self.heaps
if len(large) > len(small):
return float(large[0])
return (large[0] - small[0]) / 2.0
if __name__ == "__main__":
median_calculator = MedianCalculator()
try:
while True:
num = int(input())
median_calculator.add_num(num)
print('Current median: ', median_calculator.find_median())
except EOFError:
print('End of stream.')
@patricktrainer
Copy link
Author

patricktrainer commented Jun 15, 2023

what is this

This code uses the heapq library to maintain two heaps for a fast median calculation.

It will keep reading from stdin until an EOFError (or more specifically in this case, end of input stream) is detected.

You can test this code in a terminal by typing numbers and pressing Enter after each.

The script will print out the current median after each number. When you’re done, you can send an EOF signal with Ctrl+D

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment