15 · Sort Colors

In an Xcode test target:

import Testing

struct SortColorsTests {

    @Test(arguments: [
        ([2, 0, 1, 1, 0, 2], [0, 0, 1, 1, 2, 2]),
        ([1, 0], [0, 1]),
        ([2, 2, 1], [1, 2, 2]),
    ])
    func groupsZerosOnesAndTwosInPlace(colors: [Int], expected: [Int]) {
        var colors = colors
        sortColors(&colors)
        #expect(colors == expected)
    }

    // MARK: - Privates
    private func sortColors(_ colors: inout [Int]) {
        // your solution
    }
}

In a playground:

func sortColors(_ colors: inout [Int]) {
    // your solution
}

let cases: [([Int], [Int])] = [
    ([2, 0, 1, 1, 0, 2], [0, 0, 1, 1, 2, 2]),
    ([1, 0], [0, 1]),
    ([2, 2, 1], [1, 2, 2]),
]
for (colors, expected) in cases {
    var got = colors
    sortColors(&got)
    print(got == expected ? "PASS" : "FAIL", colors, "→", got, "expected", expected)
}

Three pointers partitioning in place — the Dutch national flag. The cue is "only three distinct values" plus "in place" and "one pass". Two pointers mark where the finished 0s end and where the finished 2s begin; a third walks the unsorted middle and throws each value to the side it belongs on.

Keep three positions. Everything before low is a finished 0. Everything after high is a finished 2. Everything from low up to mid is a finished 1, and mid to high is still unseen.

Look at the value under mid. A 0 gets swapped to low, and both low and mid move forward. A 1 is already in the right zone, so only mid moves. A 2 gets swapped to high, and high moves back — but mid stays where it is, because the value that just arrived from the far end hasn't been looked at yet. Stop when mid passes high: nothing unseen is left.

Time O(n): every step either moves mid forward or high back, so there are at most n steps. Space O(1): three indices.

func sortColors(_ colors: inout [Int]) {
    var low = 0                     // next slot for a 0
    var mid = 0                     // the value being looked at
    var high = colors.count - 1     // next slot for a 2

    while mid <= high {
        switch colors[mid] {
        case 0:
            colors.swapAt(low, mid)
            low += 1
            mid += 1
        case 1:
            mid += 1
        default:
            colors.swapAt(mid, high)
            high -= 1               // don't move mid: the value that just arrived is unseen
        }
    }
}

When a 0 is swapped from low, mid may move forward without looking, because whatever sat at low was already seen and was a 1 (or low and mid were the same position).

Verified with swift test on Swift 6.4 in Swift 6 mode: the three examples above, six more (empty, a single 0, a single 2, reversed [2, 1, 0], all 1s, [2, 0, 2, 0] with no 1s), and 3,000 random arrays checked against the library sort.

[2, 0, 1, 1, 0, 2] — low 0, mid 0, high 5

The first two rows are the reason mid holds still after a 2: the value swapped in was another 2, and it would have been left in the middle had mid moved on.

[1, 0] — the 1 at position 0 stays and mid moves to 1. The 0 there is swapped with low (position 0): [0, 1]. mid becomes 2, passes high at 1, stop.

swap(&colors[i], &colors[j]) breaks exclusive access. It is the obvious spelling and it doesn't compile: two inout accesses to the same array overlap, and Swift rejects it ("overlapping accesses to 'colors'"). Use colors.swapAt(i, j), which is the array's own method and is fine when i == j — which happens here every time low and mid sit together. The second trap is the inout itself: func sortColors(_ colors: [Int]) receives a constant copy, so an answer that "sorts in place" into a local var changes nothing the caller can see. Say "in place" and write inout and & at the call site.

• "Why not just count them?" → Counting 0s, 1s and 2s and writing them back is O(n) and O(1), but takes two passes and only works because the values carry no payload. If each entry were a record sorted by a colour key, overwriting would lose the records; swapping keeps them.
• "Is it stable?" → No. Swaps with the far end reorder equal keys. A stable version needs O(n) extra memory — three buckets appended in order.
• "Now there are k colours." → Counting sort in O(n + k), or partition repeatedly: split around the middle colour and recurse, O(n log k). The three-pointer trick doesn't generalise past three.