Master TopK: From Simple Sorts to Heap and QuickSelect Solutions

Introduction

Hello everyone, I’m Niu Niu. Today I share a classic TopK problem without considering large‑scale distributed solutions, focusing on a single algorithmic question.

TopK problem means finding the top k largest (or smallest) numbers in a sequence; its solving ideas are similar to those for the Kth largest element.

TopK appears very frequently in coding interviews. We will use LeetCode 215 (Find the Kth Largest Element in an Array) as a demonstration.

Sorting Method

Find TopK and then sort TopK.

If you use O(n²) sorting like bubble sort, you can stop after K passes because each pass fixes one maximum (or minimum) element, giving a time complexity of O(n k).

Bubble sort repeatedly swaps adjacent elements, while simple selection sort selects the maximum (or minimum) each pass and swaps it with the last unsorted element.

Process illustration:

Code for bubble/selection approach:

private void swap(int[] arr, int i, int j) {
    int temp = arr[i];
    arr[i] = arr[j];
    arr[j] = temp;
}

public int findKthLargest1(int[] nums, int k) {
    for (int i = nums.length - 1; i >= nums.length - k; i--) {
        for (int j = 0; j < i; j++) {
            if (nums[j] > nums[j + 1]) {
                swap(nums, j, j + 1);
            }
        }
    }
    return nums[nums.length - k];
}

public int findKthLargest2(int[] nums, int k) {
    for (int i = 0; i < k; i++) {
        int max = i;
        for (int j = i + 1; j < nums.length; j++) {
            if (nums[j] > nums[max]) {
                max = j;
            }
        }
        if (max != i) {
            swap(nums, i, max);
        }
    }
    return nums[k - 1];
}

Other sorts like quicksort, merge sort, or heap sort have O(n log n) complexity and can also be used, but we focus on the above two methods.

Heap‑Based Optimization

Using a max‑heap, we can build the heap in O(n) time and extract the maximum K times, each extraction costing O(log n), resulting in O(n + k log n) overall.

Heap implementation code:

class Solution {
    private void swap(int[] arr, int i, int j) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }
    private void shiftDown(int[] arr, int index, int len) {
        int left = index * 2 + 1;
        int right = index * 2 + 2;
        if (left >= len) return;
        if (right < len && arr[right] > arr[index] && arr[right] > arr[left]) {
            swap(arr, index, right);
            shiftDown(arr, right, len);
        } else if (arr[left] > arr[index]) {
            swap(arr, index, left);
            shiftDown(arr, left, len);
        }
    }
    private void createHeap(int[] arr) {
        for (int i = arr.length / 2; i >= 0; i--) {
            shiftDown(arr, i, arr.length);
        }
    }
    public int findKthLargest(int[] nums, int k) {
        createHeap(nums);
        for (int i = 0; i < k; i++) {
            int top = nums[0];
            nums[0] = nums[nums.length - 1 - i];
            nums[nums.length - 1 - i] = top;
            shiftDown(nums, 0, nums.length - i - 1);
        }
        return nums[nums.length - k];
    }
}

Quick‑Select (QuickSort) Optimization

QuickSelect uses the partition idea of quicksort to locate the Kth largest element in expected O(n) time. The algorithm recursively narrows the search range based on the position of the pivot.

QuickSelect implementation:

class Solution {
    public int findKthLargest(int[] nums, int k) {
        quickSort(nums, 0, nums.length - 1, k);
        return nums[nums.length - k];
    }
    private void quickSort(int[] nums, int start, int end, int k) {
        if (start > end) return;
        int left = start, right = end;
        int pivot = nums[start];
        while (left < right) {
            while (pivot <= nums[right] && left < right) right--;
            nums[left] = nums[right];
            while (pivot >= nums[left] && left < right) left++;
            nums[right] = nums[left];
        }
        nums[left] = pivot;
        int count = end - left + 1;
        if (count == k) return;
        if (count > k) {
            quickSort(nums, left + 1, end, k);
        } else {
            quickSort(nums, start, left - 1, k - count);
        }
    }
}

Counting Sort Extra

When the value range is limited, counting sort can find the Kth largest element in linear time by counting frequencies and scanning from the largest value.

Counting sort implementation:

public int findKthLargest(int[] nums, int k) {
    int[] cnt = new int[20001];
    int[] sum = new int[20001];
    for (int num : nums) {
        cnt[num + 10000]++;
    }
    for (int i = 20000 - 1; i >= 0; i--) {
        sum[i] = sum[i + 1] + cnt[i];
        if (sum[i] >= k) {
            return i - 10000;
        }
    }
    return 0;
}

Conclusion

TopK is not difficult; it mainly requires clever use of sorting or selection techniques. Mastering these algorithms is essential for interview success.