#include <bits/stdc++.h>
#include <gtest/gtest.h>

using namespace std;

// Given an array arr[] containing only 0s, 1s, and 2s. Sort the array
// in ascending order. Note: You need to solve this problem without
// utilizing the built-in sort function.
    
// Constraints:
// 1 ≤ arr.size() ≤ 10^6
// 0 ≤ arr[i] ≤ 2

// Expected Complexities
// Time Complexity: O(n)
// Auxiliary Space: O(1)

// Follow up: Could you come up with a one-pass algorithm using only
// constant extra space?

class Solution {
  public:
    void sort012_On(const vector<int>& values, vector<int>& arr)
    {

        // Wikipedia: The following pseudocode for three-way
        // partitioning which assumes zero-based array indexing was
        // proposed by Dijkstra himself. It uses three indices i, j
        // and k, maintaining the invariant that i ≤ j ≤ k.

        //   - Entries from 0 up to (but not including) i are values less than mid,
        //   - entries from i up to (but not including) j are values equal to mid,
        //   - entries from j up to (and including) k are values not yet sorted, and
        //   - entries from k + 1 to the end of the array are values greater than mid.        
        
        //int  low = values.at(0);
        int  mid = values.at(1);
        //int high = values.at(2);
        
        auto swap = [&arr](int x, int y)
        {
            cout << "swap idx " << x << " with idx " << y << endl;
            int tmp = arr[y];
            arr[y] = arr[x];
            arr[x] = tmp;
        };

        size_t i = 0, j = 0, k = arr.size() - 1;

        auto printIndices = [this, &arr](size_t i, size_t j, size_t k)
        {
            print(arr, string("i=" + to_string(i) +
                              " j=" + to_string(j) +
                              " k=" + to_string(k)));
        };

        // initial optimization:
        // ------------
        // while (arr[k] == high) k--;
        // while (arr[j] == low) i++, j++;

        while (j <= k)
        {
            printIndices(i, j, k);
            
            if (arr[j] < mid)
            {
                swap(i, j);
                i++; j++;
            }
            else if (arr[j] > mid)
            {
                swap(j, k);
                k--;
            }
            else //(arr[j] == mid)
            {
                j++;
            }
        }
        printIndices(i, j, k);
        print(arr, "result");
        cout << "------------------------------------" << endl;
    }
    

    // this one is O(2n)
    void sort012(vector<int>& arr)
    {
        // Index 0 num zeros, Index 1 num ones, Index 2 num twos
        vector<int> num012s{0,0,0};
    
        for (const auto& i: arr)
        {
            switch(i)
            {
                case 0: num012s.at(0)++; break;
                case 1: num012s.at(1)++; break;
                case 2: num012s.at(2)++; break;                
            }
        }

        // create sorted array
        int index = 0;
        for (int i = 0; i < num012s.at(0); i++) arr[index++] = 0;
        for (int i = 0; i < num012s.at(1); i++) arr[index++] = 1;
        for (int i = 0; i < num012s.at(2); i++) arr[index++] = 2;
        
        print(arr, "actual");
    }

    void print(const vector<int>& v, const string& tag)
    {
        cout << setw(12) << tag << ": ";
        for (const auto& j: v) cout << setw(3) << j << " ";
        cout << endl;
    }

    void buildPrefixSumArray(const vector<int>& arr)
    {
        vector<int> prefixSumArr{arr[0]};
        
        for (size_t i = 1; i < arr.size(); i++)
            prefixSumArr.push_back(prefixSumArr[i-1] + arr[i]);

        print(arr, "array");
        print(prefixSumArr, "prefixSum");
    }
};

TEST(DutchFlagAlgo, SortTest)
{
    vector<tuple<vector<int>, vector<int>, vector<int>>> v = {
        // array values  input                expected
        { {0,1,2},       {0,1,2,0,1,2},       {0,0,1,1,2,2} },
        { {0,1,2},       {2,2,1,1,0,0},       {0,0,1,1,2,2} },
        { {0,1,2},       {2,2,2,2,2,0},       {0,2,2,2,2,2} },
        { {0,1,2},       {1,1,1,1,1,0},       {0,1,1,1,1,1} },
        { {0,1,2},       {0,0,0,0,0,0,2,1},   {0,0,0,0,0,0,1,2} },
        { {0,1,2},       {2,2,2,2,0,0,0,0},   {0,0,0,0,2,2,2,2} },
        { {10,20,30},    {30,20,10,30,20,10}, {10,10,20,20,30,30} },
    };
    
    Solution s;

    for (auto i: v)
    {
        s.print(get<1>(i), "input");
        s.print(get<2>(i), "expected");

        s.sort012_On(get<0>(i), get<1>(i));

        EXPECT_EQ(get<1>(i), get<2>(i));
    }
}

int main(void)
{
    testing::InitGoogleTest();
    
    int ret = RUN_ALL_TESTS();

    // s.buildPrefixSumArray(get<0>(v.at(0)));
    
    return ret;
}