Sort Colors Solutions in C++

// Source : https://oj.leetcode.com/problems/sort-colors/
// Author : Hao Chen
// Date   : 2014-06-25



#include <stdio.h>
#include <stdlib.h>
#include <time.h>


void swap(int*a, int*b)
{
    int t;
    t=*a;
    *a = *b;
    *b = t;
}
void sortColors(int a[], int n) {
    int zero=0, two=n-1;
    for(int i=0; i<=two; i++ ){
        if (a[i]==0){
            swap(&a[zero], &a[i]);
            zero++;
        }
        if (a[i]==2){
            swap(&a[two], &a[i]);
            two--; 
            i--;
        }
    }
    
}

void printArray(int a[], int n) {
    for(int i=0; i<n; i++){
        printf("%d ", a[i]);
    }
    printf("\n");
}

int main(int argc, char** argv)
{
    int n = 7;
    if (argc>1)
        n = atoi(argv[1]);

    srand(time(NULL));

    int *a = new int[n];
    for (int i=0; i<n; i++){
        a[i] = random()%3;
    }
    printArray(a, n);
    sortColors(a, n);
    printArray(a, n);

    delete[] a;
}

#include <algorithm>

class Solution {
public:
    void sortColors(int A[], int n) {
        for (int i=0, j=0; i<n; ++i)
            if (A[i] == 0) std::swap(A[i], A[j++]);
            else if (A[i] == 2) std::swap(A[i--], A[--n]);
    }
};

/// Source : https://leetcode.com/problems/sort-colors/description/
/// Author : liuyubobobo
/// Time   : 2017-11-13

#include <iostream>
#include <vector>
#include <cassert>

using namespace std;

// Counting
// Time Complexity: O(n)
// Space Complexity: O(3)
class Solution {
public:
    void sortColors(vector<int> &nums) {

        int count[3] = {0};
        for(int i = 0 ; i < nums.size() ; i ++){
            assert(nums[i] >= 0 && nums[i] <= 2);
            count[nums[i]] ++;
        }

        int index = 0;
        for(int i = 0 ; i < count[0] ; i ++)
            nums[index++] = 0;
        for(int i = 0 ; i < count[1] ; i ++)
            nums[index++] = 1;
        for(int i = 0 ; i < count[2] ; i ++)
            nums[index++] = 2;
    }
};


void printArr(const vector<int>& vec){
    for(int e: vec)
        cout << e << " ";
    cout << endl;
}

int main() {

    vector<int> vec1 = {2, 2, 2, 1, 1, 0};
    Solution().sortColors(vec1);
    printArr(vec1);

    vector<int> vec2 = {2};
    Solution().sortColors(vec2);
    printArr(vec2);

    return 0;
}

/// Source : https://leetcode.com/problems/sort-colors/description/
/// Author : liuyubobobo
/// Time   : 2017-11-13

#include <iostream>
#include <vector>
#include <cassert>

using namespace std;

// Three Way Quick Sort
// Time Complexity: O(n)
// Space Complexity: O(1)
class Solution {
public:
    void sortColors(vector<int> &nums) {

        int zero = -1;          // [0...zero] == 0
        int two = nums.size();  // [two...n-1] == 2
        for(int i = 0 ; i < two ; ){
            if(nums[i] == 1)
                i ++;
            else if (nums[i] == 2)
                swap( nums[i] , nums[--two]);
            else{ // nums[i] == 0
                assert(nums[i] == 0);
                swap(nums[++zero] , nums[i++]);
            }
        }
    }
};


void printArr(const vector<int>& vec){
    for(int e: vec)
        cout << e << " ";
    cout << endl;
}

int main() {

    vector<int> vec1 = {2, 2, 2, 1, 1, 0};
    Solution().sortColors(vec1);
    printArr(vec1);

    vector<int> vec2 = {2};
    Solution().sortColors(vec2);
    printArr(vec2);

    return 0;
}