Algorithm to find k elements following the median in sorted order

Question

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 i^th iteration, not n.
2. Is there a better algorithm? Maybe something like O(n+klog(k))?

Answer 1

Answering your second question:

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.

Answer 2

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]

Answer 3

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.