Segregate 0’s, 1’s and 2’s in array in C++

Segregate 0's, 1's and 2's in an array

C++ Program For Sorting An Array Of 0s, 1s and 2s – When working with arrays in C++, we often come across problems where we need to rearrange or sort the elements in a specific way. One such common problem is “Segregate 0’s, 1’s, and 2’s in an array in C++.” 

This might seem like a sorting problem at first, but there’s actually a more efficient way to solve it without using any sorting algorithm.

In this blog, we’ll understand what the problem is, explore different approaches to solve it, and finally implement the most optimal solution using C++. Whether you are preparing for coding interviews or practicing DSA, this is a great example to learn how to solve problems smartly using logic and efficiency.

Segregate 0s 1s and 2s in an Array

Segregate 0’s, 1's and 2’s in an array in C++

When working with arrays, a common interview question is to segregate or sort an array that only contains the values 0, 1, and 2. This is a variation of the Dutch National Flag problem, proposed by Edsger Dijkstra.

What Does Segregating 0’s, 1’s, and 2’s Mean?

C++ Program For Sorting An Array Of 0s, 1s and 2s means rearranging the elements of the array so that all 0’s come first, then all 1’s, and finally all 2’s.

  • You are not allowed to use the sorting function and must implement the logic manually.

Example:

Input:

arr[] = {2, 0, 1, 2, 0, 1, 0}

Output:

arr[] = {0, 0, 0, 1, 1, 2, 2}
Segregate 0’s, 1's and 2’s in an array in C++-1

Algorithm for brute force

Here’s how the algorithm works:

  1. Initialize counters for 0’s, 1’s, and 2’s to 0.
  2. Loop through the array and count the occurrences of 0’s, 1’s, and 2’s.
  3. Store the counts in separate variables.
  4. Overwrite the array with the correct number of each value, starting with 0’s, then 1’s, then 2’s, using the stored counts.

Methods to Solve Segregate 0’s, 1's and 2’s in an array in C++

The task at hand is to sort an array of integers that exclusively consists of 0’s, 1’s, and 2’s, with the goal of segregating them in a specific order.

  • The desired arrangement entails placing all the 0’s before the 1’s, and all the 1’s before the 2’s.
  • This particular challenge is commonly known as the Dutch National Flag problem, drawing inspiration from the three bands of color in the Dutch flag.
  • In this blog, we will look  into two different approaches to tackle this problem.
  1. Counting Method (Brute Force)
  2. Dutch National Flag Algorithm (Optimal, 1-pass, Constant Space)
  3. STL sort() Method (Not Recommended for Interviews)
  1. Brute Force Approach(Counting Method)

Brute force approach to segregate 0’s, 1’s, and 2’s in an array is to count the number of 0’s, 1’s, and 2’s in the array, and then overwrite the array with the correct number of each value

  • Count the number of 0’s, 1’s, and 2’s
  • Overwrite the array based on these counts.

Algorithm for brute force

Here’s how the algorithm works:

  1. Initialize counters for 0’s, 1’s, and 2’s to 0.
  2. Loop through the array and count the occurrences of 0’s, 1’s, and 2’s.
  3. Store the counts in separate variables.
  4. Overwrite the array with the correct number of each value, starting with 0’s, then 1’s, then 2’s, using the stored counts.

Code for brute force method in C++

Run

#include <iostream>
using namespace std;

void segregate (int arr[], int n)
{
  int count_0 = 0, count_1 = 0, count_2 = 0;

  // Step 1: Count the number of 0's, 1's, and 2's
  for (int i = 0; i < n; i++) { if (arr[i] == 0) { count_0++; } else if (arr[i] == 1) { count_1++; } else if (arr[i] == 2) { count_2++; } } // Step 2: Overwrite the array with the correct number of each value int i = 0; while (count_0 > 0)
    {
      arr[i] = 0;
      count_0--;
      i++;
    }
  while (count_1 > 0)
    {
      arr[i] = 1;
      count_1--;
      i++;
    }
  while (count_2 > 0)
    {
      arr[i] = 2;
      count_2--;
      i++;
    }
}

int main ()
{
  int arr[] = { 2, 0, 1, 2, 0, 1, 0, 2, 1 };
  int n = sizeof (arr) / sizeof (arr[0]);

  cout << "Array: ";
  for (int i = 0; i < n; i++)
    {
      cout << arr[i] << " ";
    }
  cout << endl;

  segregate (arr, n);

  cout << "Segregated array: ";
  for (int i = 0; i < n; i++)
    {
      cout << arr[i] << " ";
    }
  cout << endl;

  return 0;
}

Output

Array: 2 0 1 2 0 1 0 2 1 
Segregated array: 0 0 0 1 1 1 2 2 2 

2. Dutch National Flag Algorithm (Optimal, 1-pass, Constant Space)

One approach to segregate 0’s, 1’s, and 2’s or C++ Program For Sorting An Array Of 0s, 1s and 2s in an array is to use the Dutch National Flag algorithm. The basic idea of this algorithm is to maintain three pointers to partition the array into three parts: the 0’s, the 1’s, and the 2’s.

  • Use 3 pointers: low, mid, and high
  • Traverse the array once and swap based on the value.

Algorithm for dutch national flag algorithm

Here’s how the algorithm works:

  1. Initialize three pointers – low pointing to the beginning of the array, mid pointing to the current element being processed, and high pointing to the end of the array.
  2. While mid is less than or equal to high, repeat steps 3-7.
  3. Check if the value at arr[mid] is 0.
  4. If yes, swap the values at arr[mid] and arr[low], and increment both mid and low.
  5. If the value at arr[mid] is 1, move mid pointer to the next element.
  6. If the value at arr[mid] is 2, swap the values at arr[mid] and arr[high], and decrement high.
  7. Repeat steps 3-6 until mid becomes greater than high.
  8. The array is now segregated with 0’s on the left, 1’s in the middle, and 2’s on the right.

Code for Dutch national flag method for Sorting An Array Of 0s, 1s and 2s in C++

Run

#include <iostream>
using namespace std;

void dutchNationalFlag (int arr[], int n)
{
  int low = 0;			// pointer to the low end of the array
  int mid = 0;			// pointer to the current element being processed
  int high = n - 1;		// pointer to the high end of the array

  while (mid <= high)
    {
      if (arr[mid] == 0)
	{
	  // Swap arr[mid] and arr[low], increment both mid and low
	  swap (arr[mid], arr[low]);
	  mid++;
	  low++;
	}
      else if (arr[mid] == 1)
	{
	  // Move mid pointer to the next element
	  mid++;
	}
      else if (arr[mid] == 2)
	{
	  // Swap arr[mid] and arr[high], decrement high
	  swap (arr[mid], arr[high]);
	  high--;
	}
    }
}

int main ()
{
  int arr[] = { 2, 0, 1, 2, 0, 1, 0, 2, 1 };
  int n = sizeof (arr) / sizeof (arr[0]);

  cout << "Array: ";
  for (int i = 0; i < n; i++)
    {
      cout << arr[i] << " ";
    }
  cout << endl;

  dutchNationalFlag (arr, n);

  cout << "Segregated array: ";
  for (int i = 0; i < n; i++)
    {
      cout << arr[i] << " ";
    }
  cout << endl;

  return 0;
}

Output

Array: 2 0 1 2 0 1 0 2 1 
Segregated array: 0 0 0 1 1 1 2 2 2 

3. STL sort() Method (Not Recommended for Interviews)

  • Use C++ STL sort from <algorithm>
  • Not efficient or acceptable in coding rounds.

Code for dutch national flag method in C++

Run
#include 
#include 
using namespace std;

int main() {
    int arr[] = {2, 0, 1, 2, 0, 1, 0};
    int n = sizeof(arr)/sizeof(arr[0]);

    sort(arr, arr + n);

    for(int i = 0; i < n; i++)
        cout << arr[i] << " ";
    cout << endl;

    return 0;
}

Output

0 0 0 1 1 2 2

Final thoughts:

C++ Program For Sorting An Array Of 0s, 1s and 2s or Segregating 0’s, 1’s, and 2’s in an array is a very common problem, especially asked in coding interviews. It helps test your understanding of arrays, loops, and sorting techniques.

There are multiple ways to solve this:

  • The Counting Method is easy to understand and implement. You just count how many 0s, 1s, and 2s are present, and then rewrite the array. But this method needs two passes – one to count, and another to update the array.
  • The Dutch National Flag Algorithm is the most efficient method. It solves the problem in just one pass and doesn’t need any extra space. It uses three pointers to place 0s, 1s, and 2s in the right position directly.
  • Using STL sort() is quick but not recommended in interviews because it doesn’t show your logical thinking and uses extra sorting power unnecessarily for such a simple case.

Get PrepInsta Prime & get Access to all 200+ courses offered by PrepInsta in One Subscription

FAQs

Because this algorithm specifically uses the fact that only three distinct elements (0, 1, 2) exist. It does the sorting in a single pass using constant space, whereas general-purpose sorting algorithms use more time (O(n log n)) and may require extra space.

No, the Dutch National Flag algorithm is specifically designed for three values. For more than three distinct values, you’ll need a different approach – often involving hashing or modified counting/sorting logic.

No, it is not a stable algorithm, meaning it does not preserve the relative order of elements with the same value. If stability is required, you’ll need a different approach, possibly involving extra space.

Yes, it performs very well even on large arrays because it’s an O(n) time and O(1) space algorithm. It’s memory efficient and avoids the overhead of sorting.

Technically, yes – but it’s not practical. Recursion would increase space complexity due to function call stack usage and could lead to stack overflow on large arrays. Iterative solutions are always better in this case.