Optimised Backtracking (Didn't works as there are multiple correct answers)


#1

A bit of maths:
S1/n1 = S2/n2 , S1+S1 = S , n1+n2 = n;
S1/n1 = (S-S1)/(n-n1)
=> S1.n-S1.n1 = S.n1-S1.n1
=> S1 = S.n1/n.
Approach :
We loop n1 : 1 to n-1 if for a given n1 S.n1%n == 0
Then we do a knapsack for S1 with each element of weight 1 and max weight n1.

#define all(x) x.begin(),x.end()
    typedef vector<vector<int>> vvi;
    typedef vector<int> vi;

bool cmp(vi a, vi b){
    if(a.size()==b.size())return a<b;
    return a.size()<b.size();
}
int n;
vi a;
vi inS1,inS2;
map<array<int,3>,pair<bool,bool>> dp;

bool find(int i, int sum, int cnt){
    if(cnt==0){
        if(sum==0){
            for(int j=i;j<n;j++)inS2.push_back(a[j]);
            return true;
        }
        return false;
    }
    if(i>=n)return false;
    if(dp[{i,sum,cnt}].first)return dp[{i,sum,cnt}].second;
    dp[{i,sum,cnt}].first = true;
    
    if(a[i]<=sum){
        inS1.push_back(a[i]);
        if(find(i+1,sum-a[i],cnt-1))return dp[{i,sum,cnt}].second = true;
        inS1.pop_back();
    }
    inS2.push_back(a[i]);
    if(find(i+1,sum, cnt))return dp[{i,sum,cnt}].second = true;
    inS2.pop_back();
    return dp[{i,sum,cnt}].second = false;
}
vector<vector<int>> Solution::avgset(vector<int> &A) {
    a.clear(); inS1.clear(); inS2.clear();
    a = A;
    n = A.size();
    int S =0;
    for(auto x: A)S+=x;
    for(int n1= 1; n1<n ; n1++){
        if((S*n1)%n==0){
            int s1 = (S*n1)/n;
            inS1.clear(); inS2.clear();
            if(find(0,s1,n1)){
                sort(all(inS1));
                sort(all(inS2));
                vvi ans = {inS1,inS2};
                sort(all(ans),cmp);
                return ans;
            }
        }
    }
    return {};
}

Any suggestions on how to find output the way the problem wants?