Advertisement
Home Data Structure Arrays Partition an Array Around a Pivot

Partition an Array Around a Pivot

0
Advertisement

Title: Partition an Array Around a Pivot

Partitioning an array around a pivot value is a core technique behind QuickSort and many selection problems. In this guide, we will explore how to rearrange an array so that every element smaller than the pivot sits on the left, every element equal to the pivot sits in the middle, and every element greater than the pivot sits on the right. You will learn the three-way partition idea, see a complete Java program, walk through a step-by-step trace, and practice interview-style follow-ups.

Understanding the Problem

Given an integer array and a pivot value, we want to reorder the elements in a single pass so the array splits into three regions: values less than the pivot, values equal to the pivot, and values greater than the pivot. The relative order of elements inside each region does not have to be preserved.

For example, if the input array is {8, 2, 6, 4, 1, 9, 5} with pivot = 5, a valid partitioned result is {2, 4, 1, 5, 6, 9, 8}. Every element on the left of 5 is smaller than it, and every element on the right is larger. The exact ordering inside each region can differ based on the swap sequence, and that is acceptable.

This operation is the single most important step in QuickSort, and it also powers the Dutch National Flag problem where the three regions are the three distinct values.

Approach to Partition the Array

We will use the three-way partition, often called the Dutch National Flag approach. It maintains three pointers that mark the boundaries of the four regions seen so far: the known-less region, the known-equal region, the unknown region, and the known-greater region.

Here are the steps:

  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] < pivot, swap arr[low] with arr[mid], then increment both low and mid.
  4. If arr[mid] > pivot, swap arr[mid] with arr[high], then decrement high. Do not advance mid because the swapped-in value has not been inspected yet.
  5. If arr[mid] equals pivot, just increment mid.
  6. Stop when mid passes high. The array is now partitioned in place.

Java Program to Partition an Array Around a Pivot

Let us put the approach into a full, compilable Java program:

Java
import java.util.Arrays;

public class PartitionArray {

    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);
                low++;
                mid++;
            } else if (arr[mid] > pivot) {
                swap(arr, mid, high);
                high--;
            } else {
                mid++;
            }
        }
    }

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

    public static void main(String[] args) {
        int[] arr = {8, 2, 6, 4, 1, 9, 5};
        int pivot = 5;
        partition(arr, pivot);

        System.out.println("Partitioned Array: " + Arrays.toString(arr));
    }
}

Explanation of the Program

  • The method partition takes the array and the pivot value as input and rearranges the array in place.
  • Pointers low and mid start at index 0, while high starts at the last index. Elements from index 0 up to low - 1 are confirmed less than the pivot; elements from low up to mid - 1 are confirmed equal to the pivot.
  • Elements from index high + 1 to the end of the array are confirmed greater than the pivot. The window mid through high is still unknown.
  • Each iteration shrinks the unknown window by at least one element, so the loop runs in O(n) time.
  • The helper swap exchanges two elements using a temporary variable, which keeps the main loop readable.
  • The main method prints the partitioned array using Arrays.toString so you can verify the result visually.

Step-by-Step Trace

Let us trace the algorithm on the input {4, 1, 7, 3, 5, 2} with pivot = 4. The three pointers start at low = 0, mid = 0, and high = 5.

Iterationarr[mid]Branch takenlow, mid, high (after)Array (after)
14equal (mid++)0, 1, 5{4, 1, 7, 3, 5, 2}
21less (swap, low++, mid++)1, 2, 5{1, 4, 7, 3, 5, 2}
37greater (swap, high–)1, 2, 4{1, 4, 2, 3, 5, 7}
42less (swap, low++, mid++)2, 3, 4{1, 2, 4, 3, 5, 7}
53less (swap, low++, mid++)3, 4, 4{1, 2, 3, 4, 5, 7}
65greater (swap, high–)3, 4, 3{1, 2, 3, 4, 5, 7}

At the end of each iteration, every element in arr[0..low-1] is less than the pivot, every element in arr[low..mid-1] equals the pivot, and every element in arr[high+1..n-1] is greater than the pivot. That invariant is the heart of why this one-pass partition is correct.

Examples

Let us look at a few concrete runs to see how the three branches interact.

Example 1: Mixed Values, No Equals

Java
int[] arr = {9, 3, 7, 1, 6, 2, 8};
int pivot = 5;
partition(arr, pivot);

System.out.println("Partitioned Array: " + Arrays.toString(arr));
// Output: Partitioned Array: [2, 3, 1, 6, 7, 8, 9]

No element equals the pivot, so the middle region stays empty. All smaller values land to the left and all larger values land to the right.

Example 2: All Elements Equal to the Pivot

Java
int[] arr = {5, 5, 5, 5};
int pivot = 5;
partition(arr, pivot);

System.out.println("Partitioned Array: " + Arrays.toString(arr));
// Output: Partitioned Array: [5, 5, 5, 5]

Every comparison takes the equal branch, so mid simply walks across the array without any swaps. The array is unchanged, which is correct because every element belongs in the middle region.

Example 3: Pivot Not Present in the Array

Java
int[] arr = {10, 20, 30};
int pivot = 15;
partition(arr, pivot);

System.out.println("Partitioned Array: " + Arrays.toString(arr));
// Output: Partitioned Array: [10, 30, 20]

The pivot does not need to be an element of the array. The algorithm still produces two regions: values below 15 on the left, values above 15 on the right.

Common Mistakes to Avoid

Here are the slip-ups learners tend to make when implementing a three-way partition:

  • Advancing mid after a greater-than swap – incrementing mid in the greater branch leaves the newly swapped-in element uninspected, so values that should move left stay on the right. Decrement high only; never touch mid in that branch.
  • Using mid < high instead of mid <= high – the loop then skips the final unknown slot, leaving an element at mid == high in the wrong region. Use <= so the last unknown position is processed.
  • Starting high at arr.length – swapping with arr[high] throws an ArrayIndexOutOfBoundsException on the first iteration. Initialize high = arr.length - 1.
  • Forgetting to swap self when low equals mid – omitting the swap in the less branch when low == mid looks harmless but makes the code inconsistent and fragile when the pivot region is non-empty. Always swap, then advance both pointers.
  • Returning a new array instead of partitioning in place – the typical interview phrasing asks for in-place work; allocating a copy inflates the space complexity from O(1) to O(n) and gives up the main benefit of this algorithm.

Edge Cases to Test

Before you declare the implementation done, run the algorithm against these boundary inputs:

  • Empty array
  • Single-element array
  • All elements equal to the pivot
  • All elements strictly less than the pivot
  • All elements strictly greater than the pivot
  • Pivot not present in the array
  • Array already partitioned around the pivot
  • Integer.MIN_VALUE or Integer.MAX_VALUE present

Additional Considerations

A few points worth remembering about this algorithm:

  1. Time complexity: The loop advances either mid or decrements high on every iteration, so each of the n elements is visited at most once. The overall running time is O(n).
  2. Space complexity: Partitioning is performed in place with only three index variables and a temporary swap slot, giving O(1) extra space.
  3. Stability: This partition is not stable – elements with the same value may be reordered because of the swaps. If you need a stable split, build two auxiliary lists and concatenate them, trading O(1) space for stability.
  4. Why Integer.MIN_VALUE matters: Even though the algorithm never does arithmetic that could overflow, picking a pivot like Integer.MIN_VALUE is a useful stress test because every element compares as greater and the greater branch runs for every iteration.

Interview Follow-up Questions

Once the base version is working, an interviewer will usually probe deeper with questions like these:

  • What if we need a stable partition? – allocate two auxiliary arrays, copy less-than elements first and greater-than elements second, then write both back; stability costs O(n) extra space.
  • How would you adapt this to partition only values 0, 1, and 2? – treat 1 as the pivot; the same three-way loop then solves the classic Dutch National Flag problem.
  • What if the array is a huge file on disk that will not fit in memory? – stream the input in chunks, write less-than and greater-than elements to two separate temporary files, then concatenate the files plus the count of equal values in the middle.
  • How would you make the partition work on a read-only input? – you cannot swap in place, so return the three regions as a tuple of lists or as index boundaries into a parallel output buffer; the time stays O(n) but the space becomes O(n).
  • Can you use this partition to find the Kth smallest element? – yes, by using it as the recursive step of QuickSelect; the expected running time becomes O(n), far faster than fully sorting the array.

Conclusion

In this guide, we walked through the three-way partition in Java step by step. The partition method rearranges the array in a single pass using three pointers, keeping the running time at O(n) and the extra space at O(1). You now have a complete program, a worked trace, concrete examples, and a checklist of common mistakes and interview follow-ups to lean on.

Try adapting the code to different pivot values, change the comparison to a comparator object, or plug the partition into QuickSort and QuickSelect to see how a tiny loop powers two of the most important algorithms in computing. Happy coding!

You can learn more about Array techniques: Arrays