quick-select algorithm to find top K numbers in an unsorted list

QuickSelect.kt

private fun MutableList<Int>.swap(a: Int, b: Int) {
    val tmp = this[a]
    this[a] = this[b]
    this[b] = tmp
}

/**
 * partitions the list between low and high and returns the size
 * of the partition.
 */
private fun MutableList<Int>.partition(low: Int, high: Int): Int {
    var size = 0
    val pivotIndex = high - 1
    val pivot = this[pivotIndex]
    for (i in low until high) {
        if (this[i] >= pivot) {
            this.swap(i, low + size)
            size++
        }
    }
    return size
}

/**
 * quick-select algorithm to find the top k items in a list
 */
private fun MutableList<Int>.topK(k: Int): List<Int> {
    check(k > 0) { "k must be positive" }

    if (k > this.size) {
        return this
    }

    var low = 0
    var high = this.size
    var remainingItems = k
    val results = ArrayList<Int>(k)
    while (remainingItems > 0) {
        val partitionSize = this.partition(low, high)
        if (partitionSize <= remainingItems) {
            results.addAll(this.drop(low).take(partitionSize))
            low += partitionSize
            remainingItems -= partitionSize
        } else {
            high = low + partitionSize - 1
        }
    }
    return results
}

fun main() {
    val numbers = MutableList(1000) { it }
    numbers.shuffle()
    val topK = numbers.topK(10)
    println(topK)
}