# Algorithm to find k elements following the median in sorted order

I have the following problem:

Given an unsorted array A of size n, print the first k elements in A larger than its median.

Here's my approach to the problem:

     1. Create a minHeap and a maxHeap
2. Iterate over elements in A // O(n)
- if maxHeap.count < minHeap.count
insert current element to maxHeap // O(log(n))
else:
insert current element to minHeap // O(log(n))
3. if maxHeap.count < minHeap.count:
median = minHeap.extractMin()
4.  output k elements from minHeap  // O(klog(n))


This maintains a maxHeap of elements less than the median and a minHeap of elements greater than or equal to the median. But from my analysis, this seems to take O(nlogn + k(log(n)) which is no better than sorting A first and grabbing A[n/2:n/2+k] in just O(nlog(n) + k).

Now I have 2 questions:

1. Is my analysis tight? I am doubtful since the heaps have at most i elements in the ith iteration, not n.
2. Is there a better algorithm? Maybe something like O(n+klog(k))?
• So please update your question with the updated algorithm, because now the code, description and the comments mismatch. I have deleted obsolete comments. Please also include the requirement that elements are sorted. – Evil Jun 28 '17 at 17:12
• How does 2. guarantee the upper half of elements to end up in minHeap? Consider k = 2, A = 6 4 2 1 3 5 7. (Actually using both heaps sound promising, but you'd need to keep max(maxHeap) no greater than mix(minHeap) - and repeated values mess up things about as good as with other approaches.) – greybeard Jun 29 '17 at 3:32
• Yeah, on second thought using heaps requires n/2 elements to be stored on both sides. My bad, I will remove the edit. – Erric Jun 29 '17 at 7:35

1. Find the median $m$, and partition the set into $L_1 = \{x | x \in S \wedge x < m\}$ and $L_2 = \{x | x \in S \wedge x > m\}$.

2. Find the $k$th largest element $x_k$ in $L_2$, and partition the set into $L_2' = \{x | x \in L_2 \wedge x \leq x_k\}$ and $L_3 = \{x | x \in L_2 \wedge x > x_k\}$.

3. Sort $L_2'$ and return.

Steps 1 and 2 take $O(n)$. Step 3 takes $O(k \log k)$. Total time $O(n + k \log k)$.

Note this assumes distinct values. Repeated values involves some edge case handling.

• @Evil, why do you need a heap? Use quick select and partition in $O(n)$. – ryan Jun 28 '17 at 21:05
• Oh, you do not. Sorry, it is due to several version. – Evil Jun 28 '17 at 21:08

The k elements following the median can be retrieved in O(n). To get them in sorted order, an additional O(k log k) sorting step is required.

### Quickselect

From a high-level perspective, the Quickselect algorithm takes an array and an integer k. It then places the k-th largest number in the correct place, smaller numbers at smaller indices and larger numbers at larger indices. The lists of smaller and larger values are unordered. This is done in O(n) on average, or even O(n) in the worst case (median of medians pivot selection).

In pseudocode a Quickselect invocation could look like this:

> arr = [Array of 101 elements in random order]
> QuickSelect(arr, 50)
> arr
[lower 50 elem's][median][upper 50 elem's]
>


### k elements following the median

Just apply Quickselect twice. On the first pass, split the input sequence into median, upper and lower half. Then apply Quickselect on the upper half to split it into a list of k-1 lower values, the k-th-order statistic and some upper values.

Here's another pseudocode implemntation:

> arr = [Array of 101 elements in random order]
> QuickSelect(arr, 50)
> arr
[lower 50 elem's][median][upper 50 elem's]
> QuickSelect(upper, k)
> arr
[lower 50 elem's][median][elem's you're looking for][median+k][upper elem's]


You can find the median of an array using QuickSelect in O(n) in the average case.

Then you only need to grab the k smallest elements out of the second half of the array in sorted order.