Sort an Array of 0, 1 and 2 (Dutch National Flag)

/*
 * Problem : Process an array.
 * Input   : n, then n integers.
 * Output  : The sorted array.
 * Approach: Three pointers.
 * Time    : O(n)
 * Space   : O(1)
 * 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) {
        // Step 1: Three zones. a[0..low-1] are all 0. a[high+1..end] are all 2. a[low..mid-1] are all 1.
        int low = 0;
        int mid = 0;
        int high = a.length - 1;

        while (mid <= high) {
            if (a[mid] == 0) {
                // Step 2: A 0 goes to the front zone.
                int temporary = a[low];
                a[low] = a[mid];
                a[mid] = temporary;
                low++;
                mid++;
            } else if (a[mid] == 1) {
                // Step 3: A 1 is already in the right zone.
                mid++;
            } else {
                // Step 4: A 2 goes to the back zone. We do not move mid, because the swapped value is not checked yet.
                int temporary = a[high];
                a[high] = a[mid];
                a[mid] = temporary;
                high--;
            }
        }
        return a;
    }

    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();
        }
        System.out.println(Arrays.stream(solve(a)).mapToObj(String::valueOf).collect(Collectors.joining(" ")));
    }
}