grouping of elements of same type together in an array in O(1) Space and O(n) time

Question

This is a problem from Elements of programming interviews book.
Question:

Given an array of n objects with keys that takes on of four values, reorder the array so that the objects that have the same type appear together. Use O(1) additional space and O(n) time.

I am unable to conceive an idea for more than three groups(Dutch National flag problem). Any help would be really appreciated.

For example a[6]={1,2,3,4,3,1}; Then the result should be {1,1,2,3,3,4} or in any order, only the grouping of numbers matters( ie, the result need not be in ascending order)

Answer 1

To make flag with same color to come together, we could have a pointer variable and swap all those elements that have the current flag to the left and repeat this for each subsequent flag. In the end, we would have all 4 flags grouped together. Space complexity is O(1) and time complexity is O(number of flags * n). Since the flags for your case is 4, it will be O(4 * n) = O(n).

#include <iostream>
using namespace std;

int main() {

    int a[] = {1,2,3,4,3,1,2,3,3,4};
    int size = sizeof(a) / sizeof(int);
    int ptr = 0;

    for(int flag=1;flag <= 4;++flag){
        for(int i=ptr;i<size;++i){
            if(a[i] == flag){
                swap(a[i],a[ptr++]);
            }
        }
    }

    for(int i=0;i<size;++i){
        cout<<a[i]<<" ";
    }

}

Update:

• For Independent of the number of groups, you can use a trie data structure. Trie would be composed of binary representation of each number with the last bit node holding the count value of occurrence for that group's element.

• Since we could represent integers in 32 bits, the space occupied by trie would be less and wouldn't grow with respect to the number of groups.

• As a result, space complexity would be O(1). Time complexity would be O(n + log(max_number) * n).