Sort Colours Problem

In the sort colours problem, we are given an array with three numbers. Let's say the numbers are 0s, 1s, and 2s. Our task is to write a program to sort these numbers. After sorting the arrays, the arrays will contain all 0s first, then 1s and at last, 2s. Let us see some examples to understand the problem:

Input: [0, 1, 0, 0, 1, 2, 0, 1, 1, 2]

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

Input: [0, 1, 0, 2, 1, 2, 1, 1, 0, 1, 1, 2, 0, 0, 1, 0, 1]

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

Brute Force Approach

The brute force approach will be to use any sorting algorithm to sort the array. The average time complexity of a good sorting algorithm is log n, where n is the size of the given array. However, we can optimize the time complexity to linear time complexity.

Approach - 1

We will solve this problem using two pointers:

The array can be divided into four sections which will store the three colours 0s, 1s, and 2s. We will initiate three pointers l = 0, m = 0 and h = last index of the array

array[0] to array[l - 1]
array[l] to array[m - 1]
array[m] to array[h - 1]
array[h] to array[n]

We will perform three operations based on the ith element:

If the current element is 0, then we will swap it with an element present at a lower index.
If the current element is 1, then we will not change anything and move forward to the next element.
If the current element is 2, then we will swap it with an element present at the higher index.

Let us see an example to understand how this program will work

array = [1, 0, 1, 2, 0, 1, 1]

l = 0, m = 0, h = 6

Step - 1:

array[m] == 1

m = m + 1 = 1

array = [1, 0, 1, 2, 0, 1, 1]

Step - 2:

array[m] = 0

swap(array[l], array[m])

m = m + 1 = 2

l = l + 1 = 1

array = [0, 1, 1, 2, 0, 1, 1]

Step - 3:

array[m] == 1

m = m + 1 = 3

array = [0, 1, 1, 2, 0, 1, 1]

Step - 4:

array[m] == 2

swap(array[m], array[h])

h = h - 1 = 5

array = [0, 1, 1, 0, 1, 1, 2]

Step - 5:

array[m] == 0

swap(array[l], array[m])

m = m + 1 = 4

l = l + 1 = 2

array = [0, 0, 1, 1, 2, 2]

Step - 6:

array[m] == 2

swap(array[m], array[h])

h = h - 1 = 4

array = [0, 0, 1, 1, 2, 2]

Step - 7:

array[m] == 2

swap(array[m], array[h])

h = h - 1 = 3

array = [0, 0, 1, 1, 2, 2]

Hence, the sorted array is array = [0, 0, 1, 1, 2, 2]

Below are the steps to be followed to solve this problem:

We will initiate three variables which will be the three-pointers. They are l = 0, m = 0, h = N.
- The complete array will be divided into four sections, one from the 0th index to l, and this section will contain 0s.
- The next section is l - m, and this range will contain 1s.
- The next section, m - h, will contain unknown numbers, and the last section, h - N, will contain 2s.
We will traverse the given array from the 0th index. The loop will continue until the middle index, i.e., m, is less than the higher index, h.
If the current element is 0, then we will swap this element with an element present at the index or pointer l. Then we will increment m and l by 1.
If the current element is 1, then we will leave the array as it is and increment the middle pointer by 1.
Lastly, if the current element is 2, we will again swap the elements. This time we will swap 2 with the element present at the higher index or at the h pointer.
- Now we will not update the m pointer as the current element can be 0, which we will need to swap with the lower element. Hence we will only update the h pointer by decrementing it by 1.
At the last of these operations, the array will be sorted and we will print the array.

Code

# Python program to sort colours or sort the given array containing 0s, 1s, and 2s

# Defining the function that will take the array and the length of that array as parameters and return the sorted array
def sort(array, n):

  # initializing the pointers
  l = 0
  h = n - 1
  m = 0
  # Iterating over the array until the condition m <= h is satisfied
  while m <= h:
    # If the element at the index m is 0
    if array[m] == 0:
      # Swap the element with the element present at the lower index
      array[l], array[m] = array[m], array[l]
      # Incrementing the pointers
      l = l + 1
      m = m + 1
    # If the element at the index m is 1
    elif array[m] == 1:
      # Do not change the array
      m = m + 1
    # If the element at the index m is 2
    else:
      array[m], array[h] = array[h], array[m]
      h = h - 1
  return array

# Defining the function to print the array
def printa(array):
  for i in array:
    print(i, end = ' ')

# Creating an array of 0s, 1s, and 2s and passing it to the function
array = [0, 1, 0, 2, 1, 2, 1, 1, 0, 1, 2, 0, 0, 1, 1, 0, 1]
n = len(array)
array = sort(array, n)
print("The sorted array is: ")
printa(array)

Output:

The sorted array is:
0 0 0 0 0 0 1 1 1 1 1 1 1 1 2 2 2

Time Complexity: Since we are using a linear loop, the time complexity of this approach is O(n).

Space Complexity: Since we have modified the original array, the space complexity of this approach is constant, i.e., O(1).

Approach - 2

In this approach, we will use the counting method to sort the array.

We will initialize three variables count0, count1, and count2.
Now, we will traverse through the given array and count the number of occurrences of each number or colour.
Lastly, we will store the colours or numbers based on the count of the numbers.

Let us see an example to understand the solution:

array = [0, 1, 0, 1, 2, 0, 1, 1, 2]

count0 = 0, count1 = 0, count2 = 0

At n = 0: array[0] == 0

count0 += 1

# count0 = 1

At n = 1: array[1] == 1

count1 += 1

# count1 = 1

At n = 2: array[2] == 0

count0 += 1

# count0 = 2

At n = 3: array[3] == 1

count1 += 1

# count1 = 2

At n = 4: array[4] == 2

count2 += 1

# count2 = 1

At n = 5: array[5] == 0

count0 += 1

# count0 = 3

At n = 6: array[6] == 1

count1 += 1

# count1 = 3

At n = 7: array[7] == 1

count1 += 1

# count1 = 4

At n = 8: array[8] == 2

count2 += 1

# count2 = 2

# Create an array containing count0 times 0, count1 times 1 and count2 times 2.

Hence, the sorted array is [0, 0, 0, 1, 1, 1, 1, 2, 2].

We will follow the following steps to solve the problem using the approach shown above.

Initializing three counters, count0 to count 0s, count1 to count the number of 1s and count2 to count the number of 2s in the array.
The next step is to find the values of count0, count1 and count2. To find the counts, we will iterate over the array and, using the if-else conditions count the number of times 0s, 1s, and 2s are present in the array.
The last step is to create an array containing the appropriate count of 0s, 1s, and 2s.

Below is the Python program for the above approach.

Code

# Python program to sort colours using the counting approach

# Defining a function that will take array of 0s, 1s, and 2s and return the sorted array
def sortarray(array, n):
  count0 = 0
  count1 = 0
  count2 = 0
 
  # Counting the number of 0s, 1s, and 2s present in the array
  for i in range(n):

    # If the current element is 0, increment count0
    if array[i] == 0:
      count0 += 1

    # If the current element is 1, increment count1
    elif array[i] == 1:
      count1 += 1

    # If the current element is 2, increment count2
    elif array[i] == 2:
      count2 += 1
 
  sortedarray = [0]*count0 + [1]*count1 + [2]*count2

  # Returning the sorted array
  return sortedarray

# Defining a function to print the array 
def printarray(array, n):
  for i in range(n):
    print(array[i], end = " ")
 
# Creating an array of 0s, 1s, and 2s and passing it to the function
array = [0, 1, 0, 2, 1, 2, 1, 1, 0, 2, 0, 0, 1, 0, 1]
n = len(array)
array = sortarray(array, n)
print("The sorted array is: ")
printarray(array, n)

Output:

The sorted array is:
0 0 0 0 0 0 1 1 1 1 1 1 2 2 2

Time Complexity: We have not used any nested loop in this program. Only linear loops are used. Therefore, the time complexity is also linear, i.e., O(N).

Space Complexity: We have created an array to store the sorted array. Hence, the space complexity is linear, i.e., O(N).

Approach - 3

We can modify the above-mentioned approach and reduce the space complexity to constant.

After finding the counts of 0s, 1s, and 2s, we will follow these steps instead of creating a new array:

Replacing the first m elements of the array with 0 where m = count0 = 3

array = [0, 0, 0, 1, 2, 0, 1, 1, 2]

Replacing the next m elements of the array with 1 where m = count1 = 4

array = [0, 0, 0, 1, 1, 1, 1, 1, 2]

Replacing the next m elements of the array with 2 where m = count2 = 2

array = [0, 0, 0, 1, 1, 1, 1, 2, 2]

Hence, the sorted array is [0, 0, 0, 1, 1, 1, 1, 2, 2].

Hence, the last step of the algorithm will be to run three while loops where the number of iterations will be equal to the value of the respective counts and store the 0s, 1s, and 2s in the original array.

Below is the Python program for the above-mentioned algorithm.

Code

# Python program to sort colours using the counting approach and sorting array in-place

# Defining a function that will take array of 0s, 1s, and 2s and return the sorted array
def sortarray(array, n):
  count0 = 0
  count1 = 0
  count2 = 0
 
  # Counting the number of 0s, 1s, and 2s present in the array
  for i in range(n):

    # If the current element is 0, increment count0
    if array[i] == 0:
      count0 += 1

    # If the current element is 1, increment count1
    elif array[i] == 1:
      count1 += 1

    # If the current element is 2, increment count2
    elif array[i] == 2:
      count2 += 1
 
  # Now that we have counted the number of 0s, 1s, and 2s, we have to update the given array according to the count.
  i = 0
 
  # Storing the count0 times 0s at the beginning of the array
  while (count0 > 0):
    array[i] = 0
    i += 1
    count0 -= 1
 
  # Now adding 1s count1 times in the middle of the array
  while (count1 > 0):
    array[i] = 1
    i += 1
    count1 -= 1
 
  # Now adding the 2s count2 times at the end of the array
  while (count2 > 0):
    array[i] = 2
    i += 1
    count2 -= 1

  # Returning the sorted array
  return array

# Defining a function to print the array 
def printarray(array, n):
  for i in range(n):
    print(array[i], end = " ")
 
# Creating an array of 0s, 1s, and 2s and passing it to the function
array = [0, 1, 0, 2, 1, 2, 1, 1, 0, 2, 0, 0, 1, 0, 1]
n = len(array)
array = sortarray(array, n)
print("The sorted array is: ")
printarray(array, n)

Output:

The sorted array is:
0 0 0 0 0 0 1 1 1 1 1 1 2 2 2

Time Complexity- Same as before, O(N).

Space Complexity: We have updated the given array; hence we have not used any extra space to store the sorted array. Therefore, the space complexity is constant, i.e., O(1).

Next TopicJump Game Problem in Python

← prev next →