Python Program For Sorting An Array Of 0s, 1s and 2s
Last Updated :
19 Sep, 2023
Given an array A[] consisting 0s, 1s and 2s. The task is to write a function that sorts the given array. The functions should put all 0s first, then all 1s and all 2s in last.
Examples:
Input: {0, 1, 2, 0, 1, 2}
Output: {0, 0, 1, 1, 2, 2}
Input: {0, 1, 1, 0, 1, 2, 1, 2, 0, 0, 0, 1}
Output: {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2}
A simple solution is discussed in this(Sort an array of 0s, 1s and 2s (Simple Counting)) post.
Method 1:
Approach:The problem is similar to our old post Segregate 0s and 1s in an array, and both of these problems are variation of famous Dutch national flag problem.
The problem was posed with three colours, here `0', `1' and `2'. The array is divided into four sections:
- a[1..Lo-1] zeroes (red)
- a[Lo..Mid-1] ones (white)
- a[Mid..Hi] unknown
- a[Hi+1..N] twos (blue)
- If the ith element is 0 then swap the element to the low range, thus shrinking the unknown range.
- Similarly, if the element is 1 then keep it as it is but shrink the unknown range.
- If the element is 2 then swap it with an element in high range.
Algorithm:
- Keep three indices low = 1, mid = 1 and high = N and there are four ranges, 1 to low (the range containing 0), low to mid (the range containing 1), mid to high (the range containing unknown elements) and high to N (the range containing 2).
- Traverse the array from start to end and mid is less than high. (Loop counter is i)
- If the element is 0 then swap the element with the element at index low and update low = low + 1 and mid = mid + 1
- If the element is 1 then update mid = mid + 1
- If the element is 2 then swap the element with the element at index high and update high = high - 1 and update i = i - 1. As the swapped element is not processed
- Print the output array.
Dry Run:
Part way through the process, some red, white and blue elements are known and are in the "right" place. The section of unknown elements, a[Mid..Hi], is shrunk by examining a[Mid]:
Examine a[Mid]. There are three possibilities:
a[Mid] is (0) red, (1) white or (2) blue.
Case (0) a[Mid] is red, swap a[Lo] and a[Mid]; Lo++; Mid++
Case (1) a[Mid] is white, Mid++
Case (2) a[Mid] is blue, swap a[Mid] and a[Hi]; Hi--
Continue until Mid>Hi.
Implementation:
Python
# Python program to sort an array with
# 0, 1 and 2 in a single pass
# Function to sort array
def sort012( a, arr_size):
lo = 0
hi = arr_size - 1
mid = 0
while mid <= hi:
if a[mid] == 0:
a[lo], a[mid] = a[mid], a[lo]
lo = lo + 1
mid = mid + 1
elif a[mid] == 1:
mid = mid + 1
else:
a[mid], a[hi] = a[hi], a[mid]
hi = hi - 1
return a
# Function to print array
def printArray( a):
for k in a:
print k,
# Driver Program
arr = [0, 1, 1, 0, 1, 2,
1, 2, 0, 0, 0, 1]
arr_size = len(arr)
arr = sort012( arr, arr_size)
print "Array after segregation :",
printArray(arr)
# This code is contributed by Harshit Agrawal
Output:
array after segregation
0 0 0 0 0 1 1 1 1 1 2 2
Complexity Analysis:
- Time Complexity: O(n).
Only one traversal of the array is needed. - Space Complexity: O(1).
No extra space is required.
Method 2:
Approach: Count the number of 0s, 1s and 2s in the given array. Then store all the 0s in the beginning followed by all the 1s then all the 2s.
Algorithm:
- Keep three counter c0 to count 0s, c1 to count 1s and c2 to count 2s
- Traverse through the array and increase the count of c0 if the element is 0,increase the count of c1 if the element is 1 and increase the count of c2 if the element is 2
- Now again traverse the array and replace first c0 elements with 0, next c1 elements with 1 and next c2 elements with 2.
Implementation:
Python
# Python implementation of the approach
# Utility function to print contents
# of an array
def printArr(arr, n):
for i in range(n):
print(arr[i],end=" ")
# Function to sort the array of 0s,
# 1s and 2s
def sortArr(arr, n):
cnt0 = 0
cnt1 = 0
cnt2 = 0
# Count the number of 0s, 1s and
# 2s in the array
for i in range(n):
if arr[i] == 0:
cnt0+=1
elif arr[i] == 1:
cnt1+=1
elif arr[i] == 2:
cnt2+=1
# Update the array
i = 0
# Store all the 0s in the
# beginning
while (cnt0 > 0):
arr[i] = 0
i+=1
cnt0-=1
# Then all the 1s
while (cnt1 > 0):
arr[i] = 1
i+=1
cnt1-=1
# Finally all the 2s
while (cnt2 > 0):
arr[i] = 2
i+=1
cnt2-=1
# Prthe sorted array
printArr(arr, n)
# Driver code
arr = [0, 1, 1, 0, 1, 2,
1, 2, 0, 0, 0, 1]
n = len(arr)
sortArr(arr, n)
# This code is contributed by shubhamsingh10
Output:
0 0 0 0 0 1 1 1 1 1 2 2
Complexity Analysis:
- Time Complexity: O(n).
Only two traversals of the array is needed. - Space Complexity: O(1).
As no extra space is required.
Please refer complete article on Sort an array of 0s, 1s and 2s for more details!
Similar Reads
Python Program for Segregate 0s and 1s in an array You are given an array of 0s and 1s in random order. Segregate 0s on left side and 1s on right side of the array. Traverse array only once. Example: Input array = [0, 1, 0, 1, 0, 0, 1, 1, 1, 0] Output array = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] Recommended: Please solve it on âPRACTICE â first, before mo
3 min read
Python3 Program for Segregate 0s and 1s in an array You are given an array of 0s and 1s in random order. Segregate 0s on left side and 1s on right side of the array. Traverse array only once. Input array = [0, 1, 0, 1, 0, 0, 1, 1, 1, 0] Output array = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1] Method 1 (Count 0s or 1s) Thanks to Naveen for suggesting this method
3 min read
Python Program For Sorting A Linked List Of 0s, 1s And 2s Given a linked list of 0s, 1s and 2s, sort it.Examples: Input: 1 -> 1 -> 2 -> 0 -> 2 -> 0 -> 1 -> NULL Output: 0 -> 0 -> 1 -> 1 -> 1 -> 2 -> 2 -> NULL Input: 1 -> 1 -> 2 -> 1 -> 0 -> NULLÂ Output: 0 -> 1 -> 1 -> 1 -> 2 -> NULL Source: Microsoft Interview | Set 1 Recommended: Please solve it on "PRAC
3 min read
Python Program For Sorting A Linked List Of 0s, 1s And 2s By Changing Links Given a linked list of 0s, 1s and 2s, sort it.Examples: Input: 2->1->2->1->1->2->0->1->0 Output: 0->0->1->1->1->1->2->2->2 The sorted Array is 0, 0, 1, 1, 1, 1, 2, 2, 2. Input: 2->1->0 Output: 0->1->2 The sorted Array is 0, 1, 2Recommended:
3 min read
Python Program to Sort an array in wave form Given an unsorted array of integers, sort the array into a wave like array. An array 'arr[0..n-1]' is sorted in wave form if arr[0] >= arr[1] <= arr[2] >= arr[3] <= arr[4] >= ..... Examples: Input: arr[] = {10, 5, 6, 3, 2, 20, 100, 80} Output: arr[] = {10, 5, 6, 2, 20, 3, 100, 80} OR
3 min read
Python Program to Sort a String Sorting strings in Python is a common and important task, whether we need to organize letters alphabetically or systematically handle text data. In this article, we will explore different methods to sort a string starting from the most efficient to the least.Using sorted with join()sorted() function
2 min read