This code is meant to find the median of a list of integers. The list is assumed to be of odd length.

This is based on this HackerRank question.

I'm doing this to improve my style and to improve my knowledge of fundamental algorithms/data structures for an upcoming coding interview.

from heapq import heappush, heappop

def median(l):
    heap = []
    for i in l:
        heappush(heap, i)
    cur_int = None
    mid = len(l) // 2
    for i in range(0, mid+1):
        cur_int = heappop(heap)
    median = cur_int
    return median

I think this algorithm is O(n) time and space complexity. Please correct me if I'm wrong. I'm also open to any suggestions to improving its complexities or general performance.


2 Answers 2

  1. This code has \$O(N log N)\$ time complexity. Building the heap by pushing elements one by one already requires \$O(N log N)\$. The heapify function works in linear time. But it doesn't matter that much, anyway, as popping \$n / 2\$ elements requires \$O(N log N)\$ time.

  2. You can achieve the same time complexity by just sorting the list. This solution is also simpler.

  3. It is possible to get \$O(N)\$ time complexity on average (and \$O(N^2)\$ in the worst case) using the following algorithm:

    • Let's choose a random pivot element and split the list into two parts: with smaller or equal elements and with bigger ones.

    • We know in which half the median is by looking at sizes of these lists.

    • Thus, we can go to the half that contains it and solve the problem recursively.

    • One can show that this algorithm has a linear time complexity in average case.

  4. There is also a deterministic algorithm for finding a median in linear time in the worst case, but it's quite complex so I will not describe it here.

  • \$\begingroup\$ Thanks for the answer. The alternatives you suggested sound like good ideas. Could you elaborate on why pushing to a heap is O(NlogN)? I was looking at the complexity table at the bottom of the wikipedia article and it says insertion in a binary heap is O(logN). \$\endgroup\$ Nov 22, 2016 at 7:00
  • \$\begingroup\$ @cycloidistic I mean O(n log n) in total. Each insertion is O (log n), but we have n of them. \$\endgroup\$
    – kraskevich
    Nov 22, 2016 at 7:10
  • \$\begingroup\$ Ah right. Makes sense. Thanks for clearing that up. \$\endgroup\$ Nov 22, 2016 at 7:13
  • \$\begingroup\$ Also, looking at that complexity table it also mentions other types of heaps. For example, it says insertion for a Fibonnaci Heap is O(1). Does that mean if we use one of those fancier heaps I could be reduced to n * O(1) = O(1)? \$\endgroup\$ Nov 22, 2016 at 7:22
  • \$\begingroup\$ @cycloidistic No. Popping n/2 elements would be n log n anyway. \$\endgroup\$
    – kraskevich
    Nov 22, 2016 at 8:55

To clean up the code a bit and increase performance, try this one line that does everything your code does :-)

from statistics import median

As far as your algorithm goes, I think that the speed would likely be improved if you just sorted the list in the first place, then selected the middle one (or ones and took the average). You don't correctly handle the case where the size of the list is even.

  • \$\begingroup\$ Thanks for the answer. My code wasn't meant to handle lists of even size since the challenge said that they would be of odd length. However, I see your point about my code not handling the even case. To do that I'd probably add a check of if the list was even or odd. If it was even then I would do what I was doing except that I would pop the 2 middle elements of the heap and return their average. Using the builtin median function makes sense but I was trying to do it myself. I probably should have added the reinventing-the-wheel tag to make it explicit. Thanks again :) \$\endgroup\$ Nov 20, 2016 at 5:01
  • 1
    \$\begingroup\$ Just added the reinventing-the-wheel tag. \$\endgroup\$ Nov 20, 2016 at 5:02

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.