Segregate 0s and 1s in an array

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,

segregate 0s and 1s in an array

The output will be as shown below.

segregate 0s and 1s in an array

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 communications@algorithmsandme.com.