Nth Element

Rearrange an array so that the element at position k is the same as in sorted order, with partition guarantees.

Nth Element

Nth Element rearranges an array so that the element at index $k$ is exactly the element that would appear at that position in a fully sorted array.

All elements before index $k$ are less than or equal to it, and all elements after are greater than or equal to it. The order within those partitions is not sorted.

This is a partition based selection primitive widely used in standard libraries.

Problem

Given an array $A$ and an index $k$, reorder the array such that:

  • $A[k]$ is the k-th smallest element
  • for all $i < k$, $A[i] \le A[k]$
  • for all $i > k$, $A[i] \ge A[k]$

Algorithm

Use a Quickselect style partition process until the pivot lands at index $k$.

nth_element(A, k):
    left = 0
    right = length(A) - 1

    while true:
        pivot_index = choose_pivot(A, left, right)
        pivot_index = partition(A, left, right, pivot_index)

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

Key Property

After execution:

$$ A[0..k-1] \le A[k] \le A[k+1..n-1] $$

The array is only partially ordered.

Example

Let

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

and

$$ k = 2 $$

After applying Nth Element, one valid result is:

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

The element at index $2$ is $4$, which matches the sorted position. The left and right sides are partitioned but not sorted internally.

Correctness

Each partition step places the pivot in its final sorted position. By restricting the search interval to the side containing index $k$, the algorithm ensures that eventually the pivot lands exactly at $k$.

The partition invariant guarantees correct ordering relative to the pivot.

Complexity

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

With randomized pivot or introselect fallback:

$$ O(n) $$

Space:

$$ O(1) $$

When to Use

Use Nth Element when:

  • you need partitioning around a rank
  • you want top-k elements without sorting
  • you plan to sort only a prefix or suffix

It is the standard primitive behind efficient partial sorting.

Implementation

import random

def nth_element(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

    left, right = 0, len(a) - 1

    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
import "math/rand"

func NthElement(a []int, k int) {
	left, right := 0, len(a)-1

	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
		}
	}
}