Problem#
Given an array with non negative numbers, divide the array into two parts such that the average of both the parts is equal. Return both parts (If exist). If there is no solution. return an empty list.
Examples#
Example 1:
1
2
3
4
5
6
7
|
Input: [1 7 15 29 11 9]
Output:
[9 15] [1 7 11 29]
The average of part is (15 + 9) / 2 = 12,
average of second part elements is (1 + 7 + 11 + 29) / 4 = 12
|
NOTE 1: If a solution exists, you should return a list of exactly 2 lists of integers A and B which follow the following condition:
- numElements in A <= numElements in B
- If numElements in A = numElements in B, then A is lexicographically smaller than B
NOTE 2: If multiple solutions exist, return the solution where length(A) is minimum. If there is still a tie, return the one where A is lexicographically smallest.
NOTE 3: Array will contain only non negative numbers.
Solution#
Code:
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
|
public static void main(String[] args) {
ArrayList<Integer> A = new ArrayList<Integer>(Arrays.asList(1, 7, 15, 29, 11, 9));
ArrayList<ArrayList<Integer>> result = avgSet(A);
System.out.println(result);
}
public static ArrayList<ArrayList<Integer>> avgSet(ArrayList<Integer> A) {
Collections.sort( A );
final ArrayList<ArrayList<Integer>> res = new ArrayList<>();
final ArrayList<Integer> path = new ArrayList<>();
int N = A.size();
int totalSum = 0;
for( int a : A ) {
totalSum += a;
}
final boolean T[][][] = new boolean[N][totalSum+1][N];
for( int i =0; i < N; i++) {
for( int j =0; j <= totalSum; j++) {
for( int k =0; k < N; k++) {
T[i][j][k] = true;
}
}
}
for( int i = 1; i < N; i++ ) {
if((totalSum*i)%N != 0) {
continue;
}
final int sumOfSet_i = (totalSum*i)/N;
if( isPossible( 0, sumOfSet_i, i, A, T, path )) {
int i1=0;
int i2=0;
final ArrayList<Integer> res2 = new ArrayList<>();
while( i2 < path.size() || i1 < N ) {
if( i2 < path.size() && path.get( i2 ) == A.get( i1 )) {
i2++;
i1++;
continue;
}
res2.add( A.get( i1++ ) );
}
if( path.size() > res2.size()) {
res.add( res2 );
res.add( path );
}
else {
res.add( path );
res.add( res2 );
}
return res;
}
}
return res;
}
private static boolean isPossible(int index, int currentSum, int len, ArrayList<Integer> A, final boolean [][][] T, final ArrayList<Integer> path) {
if( len == 0) {
return currentSum == 0;
}
if( index >= A.size()) {
return false;
}
if( !T[index][currentSum][len]) {
return false;
}
if( currentSum >= A.get( index )) {
path.add( A.get( index ) );
if( isPossible(index+1, currentSum-A.get( index ), len-1, A, T, path)) {
return true;
}
path.remove( path.size()-1 );
}
if( isPossible( index+1, currentSum, len, A, T, path )) {
return true;
}
T[index][currentSum][len] = false;
return false;
}
// Editorial
private ArrayList<Integer> array;
private boolean[][][] dp;
private ArrayList<Integer> indexSetA;
public ArrayList<ArrayList<Integer>> avgSet_Editorial(ArrayList<Integer> array) {
if (array == null || array.isEmpty()) {
return new ArrayList<>();
}
this.array = array;
Collections.sort(this.array);
int sum = 0;
for (int element : array) {
sum += element;
}
int n = array.size();
// memoization table by three states : (index, sum of setA, size of setA)
this.dp = new boolean[n][1 + sum][n];
for (int i = 0; i < n; i++) {
for (int j = 0; j < 1 + sum; j++) {
for (int k = 0; k < n; k++) {
this.dp[i][j][k] = true;
}
}
}
this.indexSetA = new ArrayList<>();
// iterate for third state : size of setA which varies from 1 to n-1
for (int sizeA = 1; sizeA < n; sizeA++) {
if ((sum * sizeA) % n != 0) {
continue;
}
int sumA = (sum * sizeA) / n;
if (isPartitionPossible(0, sumA, sizeA) == true) {
break;
}
}
return generatePartitions();
}
private boolean isPartitionPossible(final int index, final int sumA, final int sizeA) {
if (sizeA == 0) {
return sumA == 0;
}
int n = this.array.size();
if (index >= n) {
return false;
}
if (this.dp[index][sumA][sizeA] == false) {
return false;
}
if (sumA >= this.array.get(index)) {
// include the current index i.e. include the current element in setA
this.indexSetA.add(index);
if (isPartitionPossible(index + 1, sumA - this.array.get(index), sizeA - 1) == true) {
return true;
}
this.indexSetA.remove(this.indexSetA.size() - 1);
}
if (isPartitionPossible(index + 1, sumA, sizeA)) {
// skip the current index i.e. don't include the current element in setA
return true;
}
this.dp[index][sumA][sizeA] = false;
return this.dp[index][sumA][sizeA];
}
private ArrayList<ArrayList<Integer>> generatePartitions() {
int i = 0, j = 0;
int sizeA = this.indexSetA.size();
int n = this.array.size();
if (sizeA == n || sizeA == 0) { // no solution exists
return new ArrayList<>();
}
ArrayList<ArrayList<Integer>> result = new ArrayList<>();
result.add(new ArrayList<>());
result.add(new ArrayList<>());
// index i is used to iterate over all elements and index j is used to iterate over indexSetA
while (i < n && j < sizeA) {
if (i == this.indexSetA.get(j)) {
result.get(0).add(this.array.get(i));
j++;
} else {
result.get(1).add(this.array.get(i));
}
i++;
}
while (i < n) {
result.get(1).add(this.array.get(i));
i++;
}
return result;
}
|
Hint and Solution Approach#
Lets try to simplify the problem.
Lets assume the two sets are set1 and set2.
Assume sum of set1 = Sum_of_Set1, with size = size_of_set1.
Assume sum of set2 = Sum_of_Set2, with size = size_of_set2
1
2
3
4
5
6
7
8
9
10
11
12
|
SUM_of_Set1 / size_of_set1 = SUM_of_Set2 / size_of_set2
SUM_of_Set1 = SUM_of_Set2 * (size_of_set1 / size_of_set2)
total_sum = Sum_of_Set1 + Sum_of_Set2
AND size_of_set2 = total_size - size_of_set1
Sum_of_Set1 = (total_sum - Sum_of_Set1) * (size_of_set1 / (total_size - size_of_set1))
OR on simplifying,
total_sum / Sum_of_Set1 - 1 = (total_size - size_of_set1) / size_of_set1
total_sum / Sum_of_Set1 = total_size / size_of_set1
Sum_of_Set1 / size_of_set1 = total_sum / total_size
|
Note that you need the solution with minimum size_of_set1 if multiple solutions exist. So, just iterate on size_of_set1. Based on size_of_set1, you can determine the value of Sum_of_Set1. Now, the problem reduces to
1
|
Can I select size_of_set1 values from the array whose sum is Sum_of_Set1 (i.e., current_sum)?
|
Lets define our function as isPossible(ind, current_sum, current_size) which returns true if it is possible to use elements with index >= ind to construct a set of size current_size whose sum is current_sum.
1
2
3
4
5
6
7
8
9
|
isPossible(ind, current_sum, current_size : |
|
| isPossible(ind + 1, current_sum, current_size) [ Do not include current element ]
Or(|) Logical operator |
|
|
|
| isPossible(ind + 1, current_sum - value_at(ind), current_size - 1)
|
|
Can you memoize values to reduce the time complexity of the above recursive function ?