[InterviewBit] Equal Average Partition

Equal Average Partition

  • Time :
  • Space :
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
#include <bits/stdc++.h>


bool helper(vector<vector<vector<int>>>& dp, vector<int>& A, vector<int>& now, int p, int sum, int ma) {
if(ma == 0) return sum == 0;
if(p >= A.size()) return false;
if(!dp[p][sum][ma]) return false;
int& res = dp[p][sum][ma];

if(A[p] <= sum) {
now.push_back(A[p]);
if(helper(dp, A,now,p+1,sum-A[p],ma - 1)) return true;
now.pop_back();
if(helper(dp, A,now,p+1,sum,ma)) return true;
}

return res = false;
}

vector<vector<int>> Solution::avgset(vector<int> &A) {
sort(begin(A), end(A));
int sum = accumulate(begin(A), end(A), 0), n = A.size();
vector<vector<vector<int>>> dp(n + 1, vector<vector<int>>(sum + 1, vector<int>(n + 1, 1)));

for(int i = 1; i < n; i++) {
if(sum * i % n) continue;
int target = sum * i / n;
vector<int> res;
if(helper(dp, A,res,0,target,i)) {
vector<int> res2;
for(int j = 0, k = 0; k < n; k++) {
if(j < res.size() and res[j] == A[k]) j++;
else res2.push_back(A[k]);
}
return {res,res2};
}
}
return {};
}

Author: Song Hayoung
Link: https://songhayoung.github.io/2022/10/11/PS/interviewbit/equal-average-partition/
Copyright Notice: All articles in this blog are licensed under CC BY-NC-SA 4.0 unless stating additionally.