Three-Way Partition of an Array

Title: Three-Way Partition of an Array

In this guide, you will learn how to perform a three-way partition of an array around a pivot value using plain Java. The problem shows up in interviews under names like the Dutch National Flag question, and it is also the core partitioning step inside modern quicksort variants.

By the end of the article you will have a working Java program, a step-by-step trace, two complete input and output examples, and a checklist of edge cases and interview follow-up questions a real interviewer tends to ask.

Understanding the Problem

Given an integer array and a pivot value, rearrange the elements so that the array splits into three contiguous regions: elements strictly less than the pivot, elements equal to the pivot, and elements strictly greater than the pivot. The relative order inside each region is not required to be preserved, and the rearrangement should happen in place without allocating a second array.

This is a generalization of the Dutch National Flag problem introduced by Edsger Dijkstra. It runs in O(n) time and uses only a handful of index variables, so it is a strong fit when memory is tight or when the input must stay inside the same array.

Approach / Algorithm Steps

We maintain three indices: low, mid, and high. Everything to the left of low is strictly less than the pivot, everything between low and mid is equal to the pivot, and everything to the right of high is strictly greater than the pivot. The range [mid, high] is the unprocessed middle that the loop shrinks on every step.

1. Initialize low = 0, mid = 0, and high = arr.length - 1.
2. While mid <= high, inspect the element at index mid.
3. If arr[mid] is less than the pivot, swap it with arr[low], then advance both low and mid.
4. If arr[mid] is greater than the pivot, swap it with arr[high] and decrement high. Do not move mid, since the element swapped in from the right has not been inspected yet.
5. If arr[mid] is equal to the pivot, just advance mid.
6. When the loop ends the three regions are in their final positions, and no further passes are needed.

Java Program

The program below implements the algorithm as a reusable partition method and prints the rearranged array for a sample input.

public class ThreeWayPartition {
    public static void partition(int[] arr, int pivot) {
        int low = 0;
        int mid = 0;
        int high = arr.length - 1;
        while (mid <= high) {
            if (arr[mid] < pivot) {
                swap(arr, low++, mid++);
            } else if (arr[mid] > pivot) {
                swap(arr, mid, high--);
            } else {
                mid++;
            }
        }
    }

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

    public static void main(String[] args) {
        int[] arr = {2, 1, 3, 4, 2, 1, 3, 5};
        int pivot = 3;
        partition(arr, pivot);
        System.out.print("Result: ");
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i] + (i < arr.length - 1 ? ", " : ""));
        }
        System.out.println();
    }
}

Explanation of the Program

The partition method is driven by a single while loop that advances until mid crosses high. The three branches inside correspond to the three possible comparisons against the pivot, and each branch either shrinks the unprocessed window [mid, high] by one position or widens the equal region in the middle. As a result the loop always terminates after at most n iterations.

Notice why mid does not advance when the greater-than branch fires: the element that was previously at arr[high] is now at arr[mid] and has not been compared against the pivot yet. Moving mid forward in that branch would skip over an element and corrupt the partition.

The swap helper is a plain three-step exchange using a temporary variable, so the method is self-contained and does not depend on any Java collection utilities. The main method wires everything together: it builds a sample array, calls partition, and prints the elements in order using a comma separator.

Step-by-Step Trace

Let us trace the algorithm on the fresh input [4, 2, 5, 1, 3] with pivot = 3. The expected final state is [2, 1, 3, 5, 4], with the two smaller elements on the left, the single 3 in the middle, and the two larger elements on the right.

After step 5 the condition mid <= high fails, so the loop exits with the array [2, 1, 3, 5, 4]. At the end of every iteration the invariant holds: arr[0..low-1] contains elements less than the pivot, arr[low..mid-1] contains elements equal to the pivot, and arr[high+1..n-1] contains elements greater than the pivot.

Examples

Example 1: Many Duplicates of the Pivot

When the pivot occurs many times, the equal region in the middle grows quickly and the three-way partition is noticeably faster than a standard two-way split that would still have to recurse over the duplicates.

int[] arr = {0, 1, 2, 1, 0, 2, 1};
int pivot = 1;
partition(arr, pivot);
// Output: 0, 0, 1, 1, 1, 2, 2

Example 2: Already Sorted Input

A sorted input is a useful stress test because it shows that the algorithm does not preserve relative order inside a region. The greater-than branch swaps the trailing elements and leaves 4 ahead of 5 in the output.

int[] arr = {1, 2, 3, 4, 5};
int pivot = 3;
partition(arr, pivot);
// Output: 1, 2, 3, 5, 4

Common Mistakes to Avoid

A handful of specific slip-ups turn this short loop into a hard-to-debug mess. Watch for the following when writing or reviewing the code.

• Advancing mid in the greater-than branch – the element just swapped in from arr[high] has not been compared yet, so moving mid forward can leave an out-of-place greater-than element in the middle region. Keep mid fixed and let the next iteration inspect the new value.
• Using mid < high instead of mid <= high – the final element at the position where mid == high is never processed, so an input like [5, 3] with pivot 3 is left unchanged. Compare with <=, not <.
• Decrementing high without a swap – some learners write high--; alone to skip a greater-than element. This drops the element out of the partitioned regions and corrupts the result. Always swap first, then decrement.
• Caching arr[mid] across branches – after a swap the value at arr[mid] is a different element, so you cannot read it once per iteration. Access arr[mid] inside the if/else if chain so each branch sees the current value.
• Hard-coding the pivot – writing a loop that only compares against a specific literal, like 0, makes the method reusable only for the Dutch National Flag variant. Pass the pivot as a parameter so the same code handles any value.

Edge Cases to Test

Before you declare the method done, walk through the following inputs and confirm the output still obeys the three-region invariant.

• Empty array
• Single-element array
• All elements equal to the pivot
• No element equal to the pivot
• Pivot smaller than every element
• Pivot larger than every element
• Array containing Integer.MIN_VALUE or Integer.MAX_VALUE
• Null array reference

Additional Considerations

The time complexity is O(n) because each iteration either advances mid, advances both low and mid, or decrements high, so the window [mid, high] shrinks by exactly one position every step. The space complexity is O(1): aside from the three index variables and one temporary used by the swap, no auxiliary memory is allocated.

The algorithm is not stable. Elements that are equal to the pivot keep their relative order among themselves, but elements inside the less-than and greater-than regions can be reordered because swaps move them across the array. If stability matters, use an auxiliary array and copy values with a left-to-right pass instead.

The Integer.MIN_VALUE and Integer.MAX_VALUE edge cases listed above matter because the algorithm never performs arithmetic that could overflow. All comparisons are plain value comparisons, so the method stays correct even when the input straddles the extreme ends of the int range. The null-array case should either be rejected with an IllegalArgumentException or allowed through to a NullPointerException, depending on the calling convention of your project.

Interview Follow-up Questions

A real interviewer rarely stops after the base problem. Be ready for the following variants, each hinting at a deeper idea rather than a new line of code.

• What if you need a stable three-way partition? – allocate an output array and run three left-to-right scans copying the less, equal, and greater elements in order; the overall cost stays O(n) time but now uses O(n) extra space.
• How does this relate to quicksort? – the same partition is the core step of three-way quicksort, which runs in O(n log n) on average and in linear time when the input has many duplicates of the pivot.
• Can you partition around two pivots instead of one? – generalize to four regions using low, mid, and high that separate values less than, between, and greater than the two pivots; this is the partition inside Java’s own dual-pivot quicksort for primitive arrays.
• What if the array is a stream of unknown length? – the in-place algorithm no longer applies because random access is required for the high swap; instead, buffer the stream into three queues keyed by comparison against the pivot and concatenate the queues at the end.
• How do you partition in parallel across many threads? – split the input into disjoint chunks, run the three-way partition on each chunk independently, then concatenate the three regions using a prefix sum of chunk sizes to compute the final write positions.

Conclusion

The three-way partition is one of the shortest linear-time algorithms you will ever write, yet it underpins quicksort variants, the Dutch National Flag puzzle, and a surprising number of coding-interview questions. Keep the index invariants in mind, swap before you move, and the method all but writes itself.

Adapt the code to your own data: change the pivot, swap the comparison for a Comparator, or extend the method to handle objects by exchanging references instead of ints. You can learn more about array algorithms here: Arrays.