Heap selection

To keep the K largest values, use a min-heap of size K. Its root is the smallest of the K largest so far, the value a newcomer must beat. For each element: if the heap has fewer than K items, add it; otherwise, if it is larger than the root, replace the root and sift down.

Each element costs at most O(log K), so the total is O(n log K) time and O(K) space. The data may be a stream you cannot store. For the K smallest, swap in a max-heap. The animation draws the heap as a tree and marks the root: the value to beat.

1static int[] topKLargest(int[] nums, int k) {
2    PriorityQueue<Integer> heap = new PriorityQueue<>();   // min-heap
3    for (int x : nums) {
4        if (heap.size() < k) {
5            heap.add(x);
6        } else if (x > heap.peek()) {
7            heap.poll();
8            heap.add(x);
9        }
10    }
11    return heap.stream().sorted(Comparator.reverseOrder()).mapToInt(i -> i).toArray();
12}

1public static int[] topK(int[] nums, int k) {
2    PriorityQueue<Integer> heap = new PriorityQueue<>();   // min-heap
3    for (int x : nums) {
4        if (heap.size() < k) {
5            heap.add(x);
6        } else if (x > heap.peek()) {
7            heap.poll();
8            heap.add(x);
9        }
10    }
11    return heap.stream().mapToInt(i -> i).toArray();
12}

To merge K sorted lists, keep a min-heap holding the current head of each list. Pop the smallest, output it, and push the next element from the same list. With N elements in total, each costs O(log K), giving O(N log K) overall instead of O(N·K) for scanning all heads.

The animation uses a plain scan of the K heads so that the comparisons are visible; the heap version replaces that scan with a single poll(). The same pattern merges log files, search results from shards, and the runs of an external sort.

1record Node(int value, int list, int index) {}
2
3static List<Integer> mergeK(List<int[]> lists) {
4    PriorityQueue<Node> heap = new PriorityQueue<>(Comparator.comparingInt(Node::value));
5    for (int i = 0; i < lists.size(); i++) {
6        if (lists.get(i).length > 0) heap.add(new Node(lists.get(i)[0], i, 0));
7    }
8    List<Integer> out = new ArrayList<>();
9    while (!heap.isEmpty()) {
10        Node n = heap.poll();
11        out.add(n.value());
12        int next = n.index() + 1;
13        if (next < lists.get(n.list()).length) {
14            heap.add(new Node(lists.get(n.list())[next], n.list(), next));
15        }
16    }
17    return out;
18}

1public static int[] kWayMerge(int[][] runs) {
2    int k = runs.length, total = 0;
3    for (int[] r : runs) total += r.length;
4    int[] pos = new int[k];                // next unread index in each run
5    int[] out = new int[total];
6    for (int o = 0; o < total; o++) {
7        int best = -1;
8        for (int r = 0; r < k; r++) {
9            if (pos[r] < runs[r].length &&
10                (best == -1 || runs[r][pos[r]] < runs[best][pos[best]])) {
11                best = r;
12            }
13        }
14        out[o] = runs[best][pos[best]++];
15    }
16    return out;
17}

To find the k-th smallest in a matrix whose rows are sorted, or across K sorted streams, run the k-way merge but stop after k pops. You never need to look at more than the first k elements of each row, and you never need to sort the whole matrix.

Cost: O(R + k log R) for R rows, with a heap seeded with the first element of every row. If the columns are also sorted, only the first min(R, k) rows can contain one of the k smallest, so seed just those rows. For a sorted square matrix with rows and columns both sorted, binary search on the value range with a staircase count gives O(n log(max - min)).

1// k-th smallest in a matrix with sorted rows
2static int kthSmallest(int[][] m, int k) {
3    PriorityQueue<int[]> heap = new PriorityQueue<>(Comparator.comparingInt(e -> e[0]));   // {value, row, col}
4    for (int r = 0; r < m.length; r++) heap.add(new int[] { m[r][0], r, 0 });
5    int value = -1;
6    while (k-- > 0) {
7        int[] top = heap.poll();
8        value = top[0];
9        if (top[2] + 1 < m[top[1]].length) heap.add(new int[] { m[top[1]][top[2] + 1], top[1], top[2] + 1 });
10    }
11    return value;
12}

Keep the lower half of the numbers in a max-heap (low) and the upper half in a min-heap (high), with sizes equal or low one larger. Then the median is low.peek() when the count is odd and the average of the two tops when it is even.

To insert x: add it to low, move low's largest to high so every low value stays at most every high value, then rebalance if high is larger. Insertion costs O(log n) and the median is O(1). The animation shows both heaps with their tops first, plus the list of medians produced.

1class MedianFinder {
2    PriorityQueue<Integer> low  = new PriorityQueue<>(Comparator.reverseOrder()); // max-heap
3    PriorityQueue<Integer> high = new PriorityQueue<>();                          // min-heap
4
5    void add(int x) {
6        low.add(x);
7        high.add(low.poll());                 // every low value <= every high value
8        if (high.size() > low.size()) {
9            low.add(high.poll());             // keep low the same size or one larger
10        }
11    }
12
13    double median() {
14        return low.size() > high.size() ? low.peek() : (low.peek() + high.peek()) / 2.0;
15    }
16}