Find the Median - 2#4074

Find the Median - 2 (Editorial)

Let's simulate the solution to the problem using a linear array. At each step, we'll store the incoming numbers in sorted order. We find the insertion position for a new element using binary search. However, the insertion operation itself (for example, if the array is implemented using a $v ec t or$ data structure) takes linear time. Thus, the overall solution with time complexity $O (n^{2})$ O(n^2) will not fit within the time limit.

Declare two heaps: one for maintaining the minimum (a $min$ -heap), and another for maintaining the maximum (a $ma x$ -heap). We'll organize their operation so that the following conditions hold:

every element in the $ma x$ -heap is less than every element in the $min$ -heap (smaller values are stored in the $ma x$ -heap, and larger values in the $min$ -heap);
the sizes of the heaps are equal. If the total number of elements is odd, the $min$ -heap should contain one more element than the $ma x$ -heap.

When adding a new element, insert it into the $min$ -heap if it is not smaller than the minimum element of that heap; otherwise, insert it into the $ma x$ -heap. After that, perform heap rebalancing to maintain the properties described above:

If the size of the $ma x$ -heap exceeds the size of the $min$ -heap, move the largest element from the $ma x$ -heap to the $min$ -heap.
If the size of the $min$ -heap exceeds the size of the $ma x$ -heap by more than one element, move the smallest element from the $min$ -heap to the $ma x$ -heap.

The median of the current sequence is defined as follows:

if the number of elements is odd, the median is the smallest element of the $min$ -heap
if the number of elements is even, the median is equal to the arithmetic mean of the smallest element of the $min$ -heap and the largest element of the $ma x$ -heap.

Inserting an element into a heap and the subsequent rebalancing are performed in $O (l o g_{2} n)$ O(log_2n) time. Thus, the total running time of the algorithm is $O (n l o g_{2} n)$ O(nlog_2n).

Example

Let us consider the process of computing the median as numbers are received sequentially. For clarity, we'll assume that after each new number is added, the sequence is sorted. Under each sequence, the corresponding median value is shown.

After processing all elements of the example, the final state of the heaps will be as follows:

Since the number of elements is odd, the median is equal to the smallest element of the $min$ -heap, i.e. $4$ .

Let the next incoming number be $35$ . Since it is not smaller than the minimum element of the $min$ -heap, add it to the $min$ -heap. After that, the size of the $min$ -heap becomes two elements larger than the size of the $ma x$ -heap, so we need to rebalance: move the smallest element of the $min$ -heap ( $4$ ) to the $ma x$ -heap. In other words, remove $4$ from the $min$ -heap and insert it into the $ma x$ -heap.

Now the total number of elements in the array is even, so the median is defined as the arithmetic mean of the smallest element of the $min$ -heap ( $35$ ) and the largest element of the $ma x$ -heap ( $4$ ), i.e. $(35 + 4) /2 = 19$ (35 + 4) / 2 = 19. In this problem, only the integer part of the median is required to be printed.

Algorithm implementation

Declare the $min$ -heap $m M in H e a p$ mMinHeap and $ma x$ -heap $m M a x H e a p$ mMaxHeap.

priority_queue<int, vector<int>, greater<int> > mMinHeap;
priority_queue<int, vector<int>, less<int> > mMaxHeap;

Initialize the heaps. In the variable $c$ , store the number of input values received. For convenience, initialize the heaps with dummy elements.

mMaxHeap.push(-MAX); mMinHeap.push(MAX);
c = 1;

Process the next incoming number.

while(scanf("%d",&n) == 1) 
{

Depending on the value of $n$ , insert it into one of the heaps.

  if (n >= mMinHeap.top()) mMinHeap.push(n);
  else mMaxHeap.push(n);

If the size of the $ma x$ -heap exceeds the size of the $min$ -heap, or if the size of the $min$ -heap is more than one element larger than the size of the $ma x$ -heap, perform rebalancing.

  if(mMaxHeap.size() > mMinHeap.size())
  {

Move the largest element from the $ma x$ -heap to the $min$ -heap.

    mMinHeap.push(mMaxHeap.top());
    mMaxHeap.pop();
  } else

  if(mMaxHeap.size() + 1 < mMinHeap.size())
  {

Move the smallest element of the $min$ -heap into the $ma x$ -heap.

    mMaxHeap.push(mMinHeap.top());
    mMinHeap.pop();
  }

Depending on whether the number of elements in the array is even or odd, print the value of the median.

  if (c % 2) printf("%d\n",mMinHeap.top()); 
  else printf("%d\n",(mMinHeap.top() + mMaxHeap.top())/2); 
  c++;
}