Dutch National Flag / 3 Way partition problem

Recently while going through interview questions, I came across this interesting problem - Called 3-way partitioning or also called as a Dutch National flag problem.

In gist, the problem is given an array containing 3 distinct elements, (Say 0, 1 and 2) sort them in the increasing order. When you think this way, it makes sense that this is called the Dutch National flag problem (🇳🇱).

Example:

Input array

[1, 1, 2, 0, 1, 2, 0, 1, 0, 2, 1]

Output

[0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

Algorithm:

Begin with three index variables, mid, low and high
mid and low both begin with lowest index, which is 0 and high index begin with highest index which is array.size - 1
Keep the low and high index at same position in the beginning and keep moving the mid pointer and swap elements as needed
To recall, we want all zeroes to appear first, then all 1s and then all 2s at the other end
Starting with current mid position, repeat following steps from 1 to 5 until mid index is less than or equal to high index
1. If element at mid index is zero, swap it with element at low index and increment both mid and low indices
2. If element at mid index is 1, do not perform any swap and simple increment mid index, This element is situated at the position we want it to be. This is because all zeroes end up at left end, all 2s end up at right end and all 1s remain in the middle
3. If the element at mid index is 2, swap elements at mid and high indices. Decrement high index since we have already seen this element and swapped it.
It is important to note that in the step 3, we are not incrementing mid index. This is because the element which got swapped with the element at high index might be 0 and in the next iteration we may have to put it at the left end.
If we indeed do increment mid on step 3, and swapped element is 0, we will end up with an array with 0 in the middle rather than on left end
If it important to note that while swapping we could be swapping two array elements with same value. In order to avoid this, put a check in swap routine which swaps only distinct values
When mid index becomes greater than high index, stop the loop. Our array is now sorted in the increasing order of elements in the sequence 0, 1 and 2

Implementation


func threeWaySort(with array: [Int]) -> [Int] {

    // Return an array itself if its size is less then 2. It means there is nothing left to sort
    if array.count < 2 {
        return array
    }

    // Initialize all three indices.
    var low = 0, mid = 0, high = array.count - 1

    // Use temparray since we do not want to modify the original array
    var tempArray = array

    // Continue this while loop until mid <= 0 1 2 high while mid <="high" { if temparray[mid]="=" swap(num1: &temparray[low], num2: &temparray[mid]) low="low" + } else &temparray[high], - return temparray func swap( num1: inout int, int) swap only distinct elements num1="=" num2 let temp="num1" usage inputarray="[1," 1, 2, 0, 1] sortedarray="threeWaySort(with:" inputarray) code>

To understand what's going on inside an algorithm, we will print array content along with all 3 indices at the beginning of every iteration


Iteration: 0
Low: 0 Mid: 0 High: 10
[1, 1, 2, 0, 1, 2, 0, 1, 0, 2, 1]

Iteration: 1
Low: 0 Mid: 1 High: 10
[1, 1, 2, 0, 1, 2, 0, 1, 0, 2, 1]

Iteration: 2
Low: 0 Mid: 2 High: 10
[1, 1, 2, 0, 1, 2, 0, 1, 0, 2, 1]

Iteration: 3
Low: 0 Mid: 2 High: 9
[1, 1, 1, 0, 1, 2, 0, 1, 0, 2, 2]

Iteration: 4
Low: 0 Mid: 3 High: 9
[1, 1, 1, 0, 1, 2, 0, 1, 0, 2, 2]

Iteration: 5
Low: 1 Mid: 4 High: 9
[0, 1, 1, 1, 1, 2, 0, 1, 0, 2, 2]

Iteration: 6
Low: 1 Mid: 5 High: 9
[0, 1, 1, 1, 1, 2, 0, 1, 0, 2, 2]

Iteration: 7
Low: 1 Mid: 5 High: 8
[0, 1, 1, 1, 1, 2, 0, 1, 0, 2, 2]

Iteration: 8
Low: 1 Mid: 5 High: 7
[0, 1, 1, 1, 1, 0, 0, 1, 2, 2, 2]

Iteration: 9
Low: 2 Mid: 6 High: 7
[0, 0, 1, 1, 1, 1, 0, 1, 2, 2, 2]

Iteration: 10
Low: 3 Mid: 7 High: 7
[0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2]

Now, having implemented the actual code, let's write some unit tests to test the program with few boundary cases.


import XCTest
@testable import AlgorithmsPractice

class AlgorithmsPracticeTests: XCTestCase {
    
    func testDutchFlagProblem() {

        let emptyArray: [Int] = []
        let sortedArray = [0, 0, 1, 1, 2, 2]
        let reverseSortedArray = [2, 2, 1, 1, 0, 0]
        let randomArray = [0, 1, 2, 0, 2, 1, 2, 0, 1]
        let sameElementArray = [0, 0, 0, 0]
        let twoDistinctElementsArray = [0, 1, 1, 0]

        XCTAssertEqual(threeWaySort(with: emptyArray), [])
        XCTAssertEqual(threeWaySort(with: sortedArray), [0, 0, 1, 1, 2, 2])
        XCTAssertEqual(threeWaySort(with: reverseSortedArray), [0, 0, 1, 1, 2, 2])
        XCTAssertEqual(threeWaySort(with: randomArray), [0, 0, 0, 1, 1, 1, 2, 2, 2])
        XCTAssertEqual(threeWaySort(with: sameElementArray), [0, 0, 0, 0])
        XCTAssertEqual(threeWaySort(with: twoDistinctElementsArray), [0, 0, 1, 1])
    }
}

Time complexity

As you can see from above algorithm, we maintain 3 indices to traverse over an array, but in the worst case we traverse the entire array just once. So if we assume the size of input array is n, the time complexity in terms of big O notation is simply O(n)

Reference:

Dutch National Flag Algorithm - GeeksForGeeks
Dutch National Flag Algorithm - Wikipedia