Python – Sort an Array Containing Only 0, 1, and 2 without Using a Standard Sorting Algorithm

This tutorial explains how to sort an array that contains only 0s, 1s, and 2s without using any conventional sorting algorithm. We will implement the Dutch National Flag algorithm which sorts the array in a single traversal with a time complexity of O(n) and constant space, making it an excellent solution for coding interviews and for beginners learning data structures and algorithms.

Problem Statement

Given an array arr that consists solely of the numbers 0, 1, and 2, your task is to sort the array so that all 0s come first, followed by all 1s, and finally all 2s. The challenge is to accomplish this without using any standard sorting algorithm.

Sample Input and Output

Example 1:

</>
Copy
Input:  arr = [0, 2, 1, 2, 0]
Output: [0, 0, 1, 2, 2]

Example 2:

</>
Copy
Input:  arr = [2, 1, 0, 1, 2, 0, 0]
Output: [0, 0, 0, 1, 1, 2, 2]

Solution Approach

The Dutch National Flag algorithm sorts the array by maintaining three pointers:

  1. Initialize three pointers:
    • low: starting index for placing 0s.
    • mid: current index under consideration.
    • high: ending index for placing 2s.
  2. Traverse the array using the mid pointer:
    • If arr[mid] is 0, swap it with arr[low] and increment both low and mid.
    • If arr[mid] is 1, just increment mid.
    • If arr[mid] is 2, swap it with arr[high] and decrement high (do not increment mid immediately, as the swapped element needs to be checked).
  3. Repeat the process until mid is greater than high.

Python Program

</>
Copy
# Function to sort an array containing only 0s, 1s, and 2s using the Dutch National Flag algorithm
def sort_012(arr):
    low = 0                # Pointer for the next position of 0
    mid = 0                # Current element under consideration
    high = len(arr) - 1    # Pointer for the next position of 2

    while mid <= high:
        if arr[mid] == 0:
            # Swap the current element with the element at 'low' and move both pointers forward
            arr[low], arr[mid] = arr[mid], arr[low]
            low += 1
            mid += 1
        elif arr[mid] == 1:
            # Move to the next element if it is 1
            mid += 1
        else:
            # Swap the current element with the element at 'high' and move the 'high' pointer backward
            arr[mid], arr[high] = arr[high], arr[mid]
            high -= 1
    return arr

# Example usage:
# Example 1:
arr1 = [0, 2, 1, 2, 0]
print("Sorted Array:", sort_012(arr1))

# Example 2:
arr2 = [2, 1, 0, 1, 2, 0, 0]
print("Sorted Array:", sort_012(arr2))

Conclusion

In this tutorial, we explored how to sort an array containing only 0s, 1s, and 2s without using any standard sorting algorithms. The Dutch National Flag algorithm provides an optimal solution with a single traversal of the array, making it both time and space efficient.