CA 08 - Sort 0s, 1s, and 2s

Problem

Given an array arr[] containing only 0s, 1s, and 2s, the task is to sort the array in ascending order without using the built-in sort function.

Output

Example 1
Output: [0, 0, 1, 1, 2, 2]
Example 2
Output: [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2]

My Approach

To solve this problem, I used the three-pointer technique.

I maintain three pointers:

low to place 0s
mid to traverse the array
high to place 2s

Then I iterate through the array:

If the element is 0, I swap it with the element at low and move both low and mid forward
If the element is 1, I just move mid forward
If the element is 2, I swap it with the element at high and move high backward

I continue this process until mid crosses high.

This works because it places 0s at the beginning, 1s in the middle, and 2s at the end in a single traversal.

This approach is efficient because:

It requires only one traversal
It uses constant extra space
Code

def sort_array(arr):
    low = 0
    mid = 0
    high = len(arr) - 1
    while mid <= high:
        if arr[mid] == 0:
            arr[low], arr[mid] = arr[mid], arr[low]
            low += 1
            mid += 1
        elif arr[mid] == 1:
            mid += 1
        else:
            arr[mid], arr[high] = arr[high], arr[mid]
            high -= 1
    return arr