K Smallest Elements with a Heap

/*
 * Problem : K smallest.
 * Input   : n, n integers, then k.
 * Output  : The k smallest.
 * Approach: A heap of size k.
 * Time    : O(n log k)
 * Space   : O(k)
 * Java    : 8 or newer.
 */
import java.util.Scanner;
import java.util.Arrays;
import java.util.stream.Collectors;

public class Main {
    // This method solves the task. It receives the input and returns the answer.
    static int[] solve(int[] a, int k) {
        // Step 1: A max-heap keeps the k smallest values. PriorityQueue is a min-heap, so we reverse the order.
        java.util.PriorityQueue<Integer> heap = new java.util.PriorityQueue<>(java.util.Collections.reverseOrder());

        for (int value : a) {
            // Step 2: Add the value. If the heap has more than k values, remove the biggest one.
            heap.add(value);
            if (heap.size() > k) {
                heap.poll();
            }
        }

        // Step 3: The heap holds the k smallest values in no special order. Copy and sort them for the output.
        int[] result = new int[heap.size()];
        int index = 0;
        for (int value : heap) {
            result[index++] = value;
        }
        java.util.Arrays.sort(result);
        return result;
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        int size = scanner.nextInt();
        int[] a = new int[size];
        for (int i = 0; i < size; i++) {
            a[i] = scanner.nextInt();
        }
        int k = scanner.nextInt();
        System.out.println(Arrays.stream(solve(a, k)).mapToObj(String::valueOf).collect(Collectors.joining(" ")));
    }
}