# Top K by Quickselect

# Top K by Quickselect

Top K by Quickselect finds the $k$ largest elements by partitioning the array so that the top $k$ elements occupy one side. It uses the same idea as Quickselect, but returns a set instead of a single element.

This method is faster than heap-based approaches when the data is in memory and only a single pass with partitioning is needed.

## Problem

Given an array $A$ of length $n$ and an integer $k$, return the $k$ largest elements of $A$.

The output does not need to be sorted unless required.

## Algorithm

Find the $(n - k)$-th smallest element using Quickselect, then take all elements greater than or equal to it.

```id="tkq1"
top_k_by_quickselect(A, k):
    n = length(A)
    target = n - k

    nth_element(A, target)

    return A[target..n-1]
```

If sorted output is required, sort the resulting segment.

## Example

Let

$$
A = [7, 2, 9, 4, 3]
$$

and

$$
k = 3
$$

We compute:

$$
target = 5 - 3 = 2
$$

After partition:

$$
[2, 3, 4, 9, 7]
$$

Take:

$$
[4, 9, 7]
$$

These are the top 3 elements.

## Correctness

After applying Nth Element at index $n - k$, all elements to the right are greater than or equal to $A[n-k]$. Therefore they represent the $k$ largest elements.

The partition property ensures correctness without full sorting.

## Complexity

| case    | time     |
| ------- | -------- |
| average | $O(n)$   |
| worst   | $O(n^2)$ |

With randomized pivot or introselect fallback:

$$
O(n)
$$

Sorting the final $k$ elements adds:

$$
O(k \log k)
$$

Space:

$$
O(1)
$$

## When to Use

Use this method when:

* data fits in memory
* you want optimal average performance
* ordering of the top $k$ is not required or can be added later

For streaming data or strict memory bounds, heap-based methods are more appropriate.

## Implementation

```id="tkq2"
import random

def top_k_by_quickselect(a, k):
    def partition(left, right, pivot_index):
        pivot = a[pivot_index]
        a[pivot_index], a[right] = a[right], a[pivot_index]
        store = left
        for i in range(left, right):
            if a[i] < pivot:
                a[store], a[i] = a[i], a[store]
                store += 1
        a[right], a[store] = a[store], a[right]
        return store

    def nth_element(left, right, k):
        while True:
            if left == right:
                return

            pivot_index = random.randint(left, right)
            pivot_index = partition(left, right, pivot_index)

            if pivot_index == k:
                return
            elif k < pivot_index:
                right = pivot_index - 1
            else:
                left = pivot_index + 1

    n = len(a)
    target = n - k
    nth_element(0, n - 1, target)

    return a[target:]
```

```id="tkq3"
import "math/rand"

func TopKByQuickselect(a []int, k int) []int {
	var partition func([]int, int, int, int) int

	partition = func(a []int, left, right, pivotIndex int) int {
		pivot := a[pivotIndex]
		a[pivotIndex], a[right] = a[right], a[pivotIndex]

		store := left
		for i := left; i < right; i++ {
			if a[i] < pivot {
				a[store], a[i] = a[i], a[store]
				store++
			}
		}
		a[right], a[store] = a[store], a[right]
		return store
	}

	var nthElement func(int, int, int)

	nthElement = func(left, right, k int) {
		for {
			if left == right {
				return
			}

			pivotIndex := left + rand.Intn(right-left+1)
			pivotIndex = partition(a, left, right, pivotIndex)

			if pivotIndex == k {
				return
			} else if k < pivotIndex {
				right = pivotIndex - 1
			} else {
				left = pivotIndex + 1
			}
		}
	}

	n := len(a)
	target := n - k
	nthElement(0, n-1, target)

	return a[target:]
}
```