Given an array of 0s and 1s, segregate 0s and 1s in such as way that all 0s come before 1s. For example, in the array below,

The output will be as shown below.

This problem is very similar to Dutch national flag problem

## Different methods to segregate 0s and 1s in an array

**Counting 0s and 1s.**

The first method is to count the occurrence of 0s and 1s in the array and then rewrite o and 1 onto original array those many times. The complexity of this method is `O(n)`

with no added space complexity. The only drawback is that we are traversing the array twice.

package com.company; /** * Created by sangar on 9.1.19. */ public class SegregateZerosAndOnes { public void segregate(int[] a) throws IllegalArgumentException{ if(a == null) throw new IllegalArgumentException(); int zeroCount = 0; int oneCount = 0; for (int i = 0; i < a.length; i++) { if (a[i] == 0) zeroCount++; else if (a[i] == 1) oneCount++; else throw new IllegalArgumentException(); } for (int i = 0; i < zeroCount; i++) { a[i] = 0; } for (int i = zeroCount; i < zeroCount + oneCount; i++) { a[i] = 1; } } }

**Using two indices.**

the second method is to solve this problem in the same complexity, however, we will traverse the array only once. Idea is to maintain two indices, `left` which starts from index 0 and `right` which starts from end (n-1) where n is number of elements in the array.

Move `left` forward till it encounters a 1, similarly decrement `right` until a zero is encountered. If left is less than right, swap elements at these two indice and continue again.

1. Set `left` = 0 and `right` = n-1

2. While left < right
2.a if a[left] is 0 then `left++`

2.b if a[right] is 1 then `right– `;

2.c if left < right, `swap(a[left], a[right])`

### segregate 0s and 1s implementation

public void segregateOptimized(int[] a) throws IllegalArgumentException{ if(a == null) throw new IllegalArgumentException(); int left = 0; int right = a.length-1; while(left < right){ while(left < a.length && a[left] == 0) left++; while(right >= 0 && a[right] == 1) right--; if(left >= a.length || right <= 0) return; if(a[left] > 1 || a[left] < 0 || a[right] > 1 || a[right] < 0) throw new IllegalArgumentException(); if(left < right){ a[left] = 0; a[right] = 1; } } }

The complexity of this method to segregate 0s and 1s in an array is O(n) and only one traversal of the array happens.

### Test cases

package test; import com.company.SegregateZerosAndOnes; import org.junit.*; import org.junit.rules.ExpectedException; import java.util.Arrays; import static org.junit.jupiter.api.Assertions.assertEquals; /** * Created by sangar on 28.8.18. */ public class SegregateZerosAndOnesTest { SegregateZerosAndOnes tester = new SegregateZerosAndOnes(); @Test public void segregateZerosAndOnesOptimizedTest() { int[] a = {0,1,0,1,0,1}; int[] output = {0,0,0,1,1,1}; tester.segregateOptimized(a); assertEquals(Arrays.toString(output), Arrays.toString(a)); } @Test public void segregateZerosAndOnesAllZerosOptimizedTest() { int[] a = {0,0,0,0,0,0}; int[] output = {0,0,0,0,0,0}; tester.segregateOptimized(a); assertEquals(Arrays.toString(output), Arrays.toString(a)); } @Test public void segregateZerosAndOnesAllOnesOptimizedTest() { int[] a = {1,1,1,1,1}; int[] output = {1,1,1,1,1}; tester.segregateOptimized(a); assertEquals(Arrays.toString(output), Arrays.toString(a)); } @Test(expected=IllegalArgumentException.class) public void segregateZerosAndOnesOptimizedIllegalArgumentTest() { int[] a = {1,1,1,1,2}; tester.segregateOptimized(a); } @Test(expected=IllegalArgumentException.class) public void segregateZerosAndOnesOptimizedNullArrayTest() { tester.segregateOptimized(null); } }

Please share if you have any suggestion or queries. If you are interested in contributing to the website or have an interview experience to share, please contact us at [email protected]