F - Candy Redistribution Editorial by spheniscine

Solution in O(N * 2^N)

Let \(\displaystyle m = \frac {\sum _i A_i} N\) represent the arithmetic mean of the candies over the children. If this is not an integer, it’s obvious that there’s no solution.

Otherwise there is a solution. Since we want all children to have exactly \(m\) candies at the end, let \(B_i = A_i - m\) represent the “balance” of candies of a child, and for any subset of children \(S\), let the balance of the subset be the sum of the balances of the children in it: \(\mathrm{bal}(S) = \sum_{i \in S} B_i\)

For any subset with \(n\) children and a balance of \(0\), we can make them all have \(m\) candies using at most \(n-1\) operations:

  • designate an arbitrary child as the “hub”
  • for every child other than the hub with a positive balance, have them give their surplus candy to the hub
  • for every child other than the hub with a negative balance, have the hub give them the candies they are short of

So now the problem boils down to finding a maximal partition into disjoint zero-balance subsets; this can be done using bitmask DP. Let \(\mathrm{dp}(S)\) be the cost of a subset, it can be seen that \(\mathrm{dp}(S) = \min _{i \in S} \mathrm{dp}(S \setminus i) + [\mathrm{bal}(S \setminus i) \ne 0]\), where square brackets evaluate to \(1\) if the expression therein is true and \(0\) otherwise. To reconstruct the solution, you can update masks of the next zero-balance subset in the partition whenever you update the DP.

This solution should take \(O(N \cdot 2^N)\) time and \(O(2^N)\) space.

posted:
last update: