Recently while going through interview questions, I came across this interesting problem  Called 3way 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
andhigh

mid
andlow
both begin with lowest index, which is0
and high index begin with highest index which isarray.size  1

Keep the
low
andhigh
index at same position in the beginning and keep moving themid
pointer and swap elements as needed 
To recall, we want all
zeroes
to appear first, then all1s
and then all2s
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 If element at
mid
index is zero, swap it with element atlow
index and increment bothmid
andlow
indices  If element at
mid
index is 1, do not perform any swap and simple incrementmid
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  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.
 If element at

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 athigh
index might be0
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 thanhigh
index, stop the loop. Our array is now sorted in the increasing order of elements in the sequence0
,1
and2
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 <= high
while mid <= high {
if tempArray[mid] == 0 {
swap(num1: &tempArray[low], num2: &tempArray[mid])
low = low + 1
mid = mid + 1
} else if tempArray[mid] == 1 {
mid = mid + 1
} else if tempArray[mid] == 2 {
swap(num1: &tempArray[high], num2: &tempArray[mid])
high = high  1
}
}
return tempArray
}
func swap( num1: inout Int, num2: inout Int) {
// Swap only distinct elements
if num1 == num2 {
return
}
let temp = num1
num1 = num2
num2 = temp
}
// Usage
let inputArray = [1, 1, 2, 0, 1, 2, 0, 1, 0, 2, 1]
let sortedArray = threeWaySort(with: inputArray)
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