https://www.interviewbit.com/problems/equal-average-partition/

AveragePartitionning.cpp

/* Copyright 2017 Rafael Rendón Pablo <rafaelrendonpablo@gmail.com> */
// region Template
#include <iostream>
#include <iomanip>
#include <fstream>
#include <sstream>
#include <string>
#include <bitset>
#include <vector>
#include <queue>
#include <map>
#include <set>
#include <cmath>
#include <cctype>
#include <cassert>
#include <cstring>
#include <numeric>
#include <iomanip>
#include <algorithm>

using std::string;
using std::vector;
using std::map;
using std::set;
using std::queue;

typedef long long           int64;
typedef unsigned long long  uint64;
// endregion

class Solution {
public:
    vector<vector<int>> avgset(vector<int> &A);
};


const int kMaxN = 80;
const int kMaxSum = 650;
const short kInf = 1 << 14;

struct Node {
    short int pos, sum, count;
    Node() {
        pos = -1;
        sum = -1;
        count = -1;
    }

    Node(int pos, int sum, int count) {
        this->pos = pos;
        this->sum = sum;
        this->count = count;
    }

    bool operator!=(const Node& that) const {
        return (this->pos != that.pos) || (this->sum != that.sum) || (this->count != that.count);
    }
    
    bool operator==(const Node& that) const {
        return this->pos == that.pos && this->sum == that.sum && this->count == that.count;
    }

    void print() {
        std::cout << "pos = " << pos
            << " sum = " << sum
            << " count = " << count << std::endl;
    }
};


Node kNotFound{kInf - 1, kInf - 1, kInf - 1};
Node kEndOfSolution{kInf, kInf, kInf};

Node dp[kMaxN][kMaxSum+1][kMaxN];

vector<int> findSolution(const int k, const int sum, const int size, const vector<int>& A) {
    // memset(dp, -1, sizeof dp);
    vector<int> path;
    for (int p = size; p >= 0; p--) {
        for (int s = kMaxSum; s >= 0; s--) {
            for (int c = k; c >= 0; c--) {
                Node& ans = dp[p][s][c];
                if (p >= size) {
                    if (c != k) {
                        ans = kNotFound;
                        continue;
                    }
                    if (s == 0) {
                        if (sum == 0) {
                            ans = kEndOfSolution;
                        } else {
                            ans = kNotFound;
                        }
                        continue;
                    }

                    int x = sum / std::__gcd(sum, size);
                    int y = size / std::__gcd(sum, size);

                    int p = s / std::__gcd(s, c);
                    int q = c / std::__gcd(s, c);
                    if (x == p && y == q) {
                        ans = kEndOfSolution;
                    } else {
                        ans = kNotFound;
                    }
                    continue;
                }

                if (c < k && (s + A[p]) < kMaxSum) {
                    Node next = dp[p+1][s+A[p]][c+1];
                    if (next != kNotFound) {
                        ans = Node{p + 1, s + A[p], c + 1};
                        continue;
                    }
                }

                Node skip = dp[p+1][s][c];
                if (skip != kNotFound) {
                    ans = Node{p + 1, s, c};
                    continue;
                }

                ans = kNotFound;
            }
        }
    }
    return dp[0][0][0];
}

vector<vector<int>> Solution::avgset(vector<int> &A) {
    std::sort(A.begin(), A.end());
    vector<vector<int>> solutions;
    int sum = std::accumulate(A.begin(), A.end(), 0);
    int size = A.size();

    for (int c = 1; c < size; c++) {
        if (findSolution(c, sum, size, A) == kNotFound) {
            continue;
        }

        set<int> x;
        Node node{0, 0, 0};
        while (node != kEndOfSolution) {
            int pos = node.pos;
            int sum = node.sum;
            int count = node.count;
            if (dp[pos][sum][count].count > count) {
                x.insert(pos);
            }
            node = dp[pos][sum][count];
        }
        vector<int> first, second;
        for (int i = 0; i < size; i++) {
            if (x.find(i) != x.end()) {
                first.push_back(A[i]);
            } else {
                second.push_back(A[i]);
            }
        }
        solutions.push_back(first);
        solutions.push_back(second);

        if (solutions[0].size() > solutions[1].size()) {
            std::swap(solutions[0], solutions[1]);
        } else if (solutions[0].size() == solutions[1].size()) {
            std::sort(solutions.begin(), solutions.end());
        }
        return solutions;
    }
    return solutions;
}


int main() {
    Solution solution;
    vector<int> A{1, 2, 3, 4, 5, 6};
    // vector<int> A{1, 7, 15, 29, 11, 9};
    // vector<int> A{24,13,30,9,13,17,42,28,20,9,24,14,20,1,28,24,26,45,27,34,38,48,38,7,34};
    // vector<int> A{19, 5, 38, 22, 44, 12, 17, 35};
    // vector<int> A{18, 29, 0, 47, 0, 41, 40, 28, 7, 1};
    auto sol = solution.avgset(A);
    for (auto v : sol) {
        for (int x : v) {
            std::cout << x << " ";
        }
        std::cout << std::endl;
    }
    return EXIT_SUCCESS;
}