F - Candy Redistribution Editorial by en_translator
Observations
The objective is infeasible if \(A_1 + \dots + A_N\) is not a multiple of \(N\), and feasible if it is a multiple of \(N\).
Let \(X := \displaystyle \frac{A_1 + \dots + A_N}{N}\). The condition in the problem statement essentially asks to make \(a_1, \dots, a_N\) all equal to \(X\). (In this editorial, \(a_1, \dots, a_N\) will denote the sequence to be operated, and \(A_1, \dots, A_N\) will denote the input initial values.)
Suppose that a sequence of operations \((x_1, y_1, z_1), \dots, (x_q, y_q, z_q)\) satisfies the condition. Consider an \(N\)-vertex graph with edges between \(q\) vertex pairs \((x_1, y_1), \dots, (x_q, y_q)\), and let \(V_1, \dots, V_s\) be its connected components. Then each \(V_k\) satisfies:
\[\sum_{v \in V_k} A_v = X \times |V_k|.\]
(The sum of \(a_v\) remains invariant by the operations, and \(a_v\) all end up being \(X\) after the operations completed.) Also, the connected component \(V_k\) has at least \((|V_k| - 1)\) edges, so the whole graph has at least \((N - s)\) edges.
Conversely, given a partition \(V_1 \sqcup \dots \sqcup V_S = \{1, \dots, N\}\) of the children such that each subset \(V_k\) satisfies \(\sum_{v \in V_k} A_v = X \times |V_k|\), one can construct an operation sequence of length \((N-s)\) that satisfies the condition in the problem statement in the following way:
- Let \(t := |V_k|\) be the number of children in each subset \(V_k\), and \(v_{k, 1}, \dots, v_{k, t}\) be its elements in descending order of \(A_{v_{k,i}}\). For each \(i = 1, \dots, t-1\) in ascending order, perform an operation \((x,y,z)\), where \((x,y,z)\) is defined as:
- \(x := v_{k,i}\), \(y := v_{k,i+1}\), and \(z := a_{x} - X\). Here, \(z \geq 0\) is guaranteed because \(v_{k,i}\) is in descending order.
Therefore, it is sufficient to find a way to partition the multiset \(A = \{A_1, \dots, A_N\}\) into subsets whose average is \(X\) so that the number \(s\) of subsets are maximized, and perform the operations above for each subset.
- One can find such a partition with a bit DP (Dynamic Programming) in which elements are inserted one by one, in a total of \(O(N 2^N)\) time.
- One can obtain the operation sequence in \(O(N \log N)\) time by implementing the algorithm above.
Hence, the problem can be solved in \(O(N 2^N)\) time, which is fast enough.
Sample code (C++)
#include <iostream>
using std::cin;
using std::cout;
using std::min;
using std::max;
#include <vector>
using std::vector;
using std::pair;
#include <algorithm>
using std::sort;
typedef long long int ll;
ll n;
vector<ll> a;
vector<ll> dp, subsum, dprev;
vector<ll> x, y, z;
void solve () {
ll sum = 0;
for (ll i = 0; i < n; i++) {
sum += a[i];
}
if (sum % n != 0) {
cout << "-1" << "\n";
return;
}
ll avg = sum / n;
//
dp.resize(1LL << n);
subsum.resize(1LL << n);
dprev.resize(1LL << n);
dp[0] = 0;
subsum[0] = 0;
dprev[0] = -1;
for (ll b = 1; !(b >> n); b++) {
dp[b] = n+1;
dprev[b] = -1;
subsum[b] = 0;
for (ll i = 0; i < n; i++) {
if (b & (1LL << i)) subsum[b] += (a[i] - avg);
}
ll item = ((subsum[b] == 0) ? 0 : +1);
for (ll i = 0; i < n; i++) {
if (b & (1LL << i)) {
ll maybe = dp[b - (1LL << i)] + item;
if (dp[b] > maybe) {
dp[b] = maybe;
dprev[b] = i;
}
}
}
}
ll cb = (1LL << n) - 1;
while (cb) {
vector<pair<ll, ll> > v;
while (true) {
ll vi = dprev[cb];
v.push_back({-a[vi], vi}); // a[i] in decreasing order
cb -= (1LL << vi);
if (subsum[cb] == 0) break;
}
sort(v.begin(), v.end());
ll carry = a[v[0].second] - avg;
for (ll i = 1; i < v.size(); i++) {
x.push_back(v[i-1].second);
y.push_back(v[i-0].second);
z.push_back(carry);
carry += (a[v[i].second] - avg);
}
}
cout << z.size() << "\n";
for (ll i = 0; i < z.size(); i++) {
cout << (x[i] + 1) << " " << (y[i] + 1) << " " << z[i] << "\n";
}
}
int main (void) {
cin >> n;
a.resize(n);
for (ll i = 0; i < n; i++) {
cin >> a[i];
}
solve();
return 0;
}
posted:
last update: