Three way partitioning of an array around a given value

The problem, simply

You get an array and a range [a, b]. Rearrange the array into three groups, in this order:

1. numbers smaller than a,
2. numbers from a to b (both ends included),
3. numbers bigger than b.

The order inside each group does not matter. Do it in one pass and without a second array.

A story to remember it

Picture it

Examples

Way 1: the simple way

Use three bags. Walk the array and drop each number in the small, medium or big bag. Then pour the bags back in that order.

class Solution {
    public void threeWayPartition(int[] arr, int a, int b) {
        List<Integer> small = new ArrayList<>(), mid = new ArrayList<>(), big = new ArrayList<>();
        for (int x : arr) {
            if (x < a) small.add(x);
            else if (x > b) big.add(x);
            else mid.add(x);
        }
        int k = 0;
        for (int x : small) arr[k++] = x;
        for (int x : mid) arr[k++] = x;
        for (int x : big) arr[k++] = x;
    }
}

How long it takes: two walks, O(n). Extra room: the bags hold all n numbers, O(n). Sorting the whole array would also work, but it costs O(n log n) and does more than needed.

Way 2: the clever way

Keep three fingers:

• low: everything before it is small.
• mid: the number we look at now. Everything from low to mid - 1 is medium.
• high: everything after it is big. Everything from mid to high is not looked at yet.

While mid <= high:

1. If arr[mid] < a: swap it with arr[low], then move both low and mid right.
2. If arr[mid] > b: swap it with arr[high] and move high left. Do not move mid: the number that came back is new.
3. Otherwise it is medium: move mid right.

Green is the small zone, blue the medium zone, red the big zone.

class Solution {
    public void threeWayPartition(int[] arr, int a, int b) {
        int low = 0, mid = 0, high = arr.length - 1;
        while (mid <= high) {
            if (arr[mid] < a) {
                swap(arr, low, mid);     // into the small zone
                low++;
                mid++;
            } else if (arr[mid] > b) {
                swap(arr, mid, high);    // into the big zone
                high--;                  // mid stays: the number that came back is new
            } else {
                mid++;                   // medium: already in the right zone
            }
        }
    }

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

How long it takes: each step moves mid right or high left, so at most n steps, O(n). Extra room: three fingers, O(1).

Watch out for

• After a swap with low, it is safe to move mid: the number that comes back from low was already checked and is medium (or it is the same number).
• The loop test is mid <= high, not mid < high. The number at high itself is not checked yet.
• The range ends are inside: use < a and > b, not <= and >=.

What interviewers ask next

• Sort an array of 0s, 1s and 2s (LeetCode 75). This is the same loop with a = b = 1.
• Can you keep the order inside each group? Not with swaps in one pass. Use the three bags, which is a stable partition.
• Why does quicksort use a three-way partition? When many numbers equal the pivot, putting them all in the middle group means they are never sorted again. Two-way quicksort can slow down to O(n^2) on many repeats.
• Return where each group starts. After the loop, the small group is 0..low-1, the medium group is low..high, and the big group is high+1..n-1.

Remember this

Grown-up words: the three zones are a three-way partition, this loop is Dijkstra's Dutch national flag algorithm, and the zones staying correct at every step is the loop invariant.