← All solutions

Shortest judge solution: 2510 bytes. Shortest team solution (during contest): N/A bytes.

This was one of the two very hard problems in the set. What is not hard to see is that the problem is simply asking for a minimum cost maximum flow between a set of sources and a set of sinks in a flow graph. The flow graph in question has a special structure: most importantly, it is a tree (and in addition, all edge capacities are infinite – we can move as many armies as we like across an edge – only the source and sink vertices have limited capacities). However the tree can have around 250 000 vertices/edges, so running a general min cost max flow algorithm on this would time out rather badly. So to solve the problem we have to (unsurprisingly) exploit the structure of the graph.

Let us without loss of generality make the following two assumptions:

1. There is no demand/supply on internal vertices of the input, or at the root in case that is also a leaf (for every internal node we can always connect it with an edge of cost 0 to a newly created child leaf which has the demand/supply).
2. The tree is binary (for every internal node with degree d ≥ 3, we can make a chain of d − 1 internal nodes connected by edges of cost 0).

(Applying the transformations may blow up the size of the tree by a factor 4, so from an implementation point of view, one might want to unravel what the transformations mean in the context of the solution below and not actually do the transformations on the tree – the solution works without these assumptions an we have made them only to simplify the presentation.)

Root the tree arbitrarily. A natural idea is to try to solve this by some form of dynamic programming on the tree, moving upwards from the leaves until we reach the root where we get the answer. But even so, it is not at all clear what state to keep in the dynamic program, or how to go up the tree efficiently.

Let us first describe the state to use, and how to use it, without worrying too much about efficiency, and then we will later figure out how to actually make it efficient. For a rooted tree T (which you can think of for now as the subtree of the input below some vertex v), define a function c_T : ℤ → ℤ, where c_T(a) is the answer to the question “What is the minimum cost of movements within the tree T, assuming that exactly a new armies have been added to the root of T?” If we have this function computed for T the entire input tree, the answer to the problem is then simply c_T(0). Note that c_T is defined also for negative a, which is interpreted as an additional demand of a armies having been introduced at the root. Let ∆_T denote the net demand within T (sum of y_i values minus sum of x_i values). Then c_T(a) is a well-defined non-increasing function for a ≥ ∆_T.

We now use this setup to make a solution that is too slow, running in time Θ(nX) where X is the sum of all x_i values in the graph. The idea is as mentioned above to build the functions c_T for subtrees of the input starting at the leaves and moving upwards. For T being a single leaf node of the input, the function c_T is trivial: c_T(a) = 0 for all a (there can never be any costs from movement within a tree of a single node). In order to progress up the tree, we need to be able to do the following two operations:

1. Extend a tree T rooted at some vertex v upwards by adding the parent w of v from the input, and joining v and w so that w becomes the new root. Call this new tree T₀. Then, ∆_T₀ = ∆_T and c_T(a) = min(c_T(a − 1), c_T₀(a) + |a| cost(v, w)). The min here is because when adding a new armies to the root we have two options: either we ignore at least one army (and get cost c_T(a − 1)) or we transport all of them to v (and get cost c_T₀(a − 1) + |a| cost(v, w)). For negative a one can see that the min is always achieved by the second option (which is good because in this setting we are moving some armies out of the tree, and do not have the option of ignoring any of them).
2. Join the the two subtrees of a vertex v. Suppose T_L and T_R are the left and right subtrees of v after doing the extend operation above (so that they are both rooted at v), and let T denote the entire subtree rooted at v. Then it is not hard to see that ∆_T = ∆_{T_L} + ∆_{T_R} and c_T(a) = min_{∆T_L ≤ a0 ≤ a − ∆T_R} c_{T_L}(a₀) + c_{T_R}(a − a₀).

To get the slow Ω(nX) solution, we can simply represent the functions c_T as arrays of function values and implement the equations written above.

In order to speed this up, the perhaps key observation is that the functions c_T are not only non-increasing, they are also convex (proof omitted, it can be derived from the identies above, or more abstractly from properties of min cost max flows). I.e., there is a diminishing returns type of property – the savings in cost attained by moving in more armies decreases as we move in more armies. For this reason, instead of working with c_T directly, it turns out to be more convenient to work with the function f_T : ℤ_≥0 → ℤ defined by f_T(a) = c_T(∆_T + a) − c_T(∆_T + a + 1) (i.e., f_T(0) says how much the cost decreases when going from ∆_T to ∆_T + 1 armies moving out of T, and so on). So we will keep this function, and the value b_T := c_T(∆_T), the base cost for the tree T if we only keep the bare minimum of soldiers inside it and move everyone else out of it. The fact that c_T is non-increasing and convex then translates into the function f_T being non-negative and non-increasing.

How does this change of perspective affect the identites for c_T above? The join operation becomes particularly nice in this new perspective: it follows from the convexity property that if we take all the f_{T_L} and f_{T_R} values and merge them into a big list sorted in non-increasing order, then that resulting list are the values of f_T (abstractly, this is because the operation we are doing is taking the Minkowski sum of the two functions c_{T_L} and c_{T_R}, and the Minkowski sum of convex sets behaves nicely). This suggest the following way of representing the functions f_T: keep the values as a sorted set, and when joining f_{T_L} with f_{T_R}, add the values from the smaller of the two sets to the larger of the two (and be careful not to do anything which takes linear time in the larger of the two sets). By a standard argument, this results in a total running time of O(n + X log₂ X) for doing all the joins.

For the extend operation, the identity for c_T above now becomes f_T₀(a) = max(0, f_T(a) − sgn(∆_T + a) cost(v, w)) (where we define sgn(a) = +1 if a ≥ 0 and −1 if a < 0). Computing this by updating each individual f_T-value would be too costly. Ignoring for a moment the cap at 0 and focusing only on the second argument, we see that the transformation performed on f_T is fairly simple: it gets increased by cost(v, w) for all a < −∆_T (which are none, if ∆_T ≥ 0, i.e., if the subtree does not have more surplus than demand), and decreased by the same amount for all a ≥ −∆_T.

We can arrange for this by augmenting the set structure suggested above with a split point and two global shifts – everything below the split point has been shifted by the first shift, and everything above the split point has been shifted by the second split. When we do the extend operation, there may already be an existing split point and shifts, we then have to move the split point, which can be done in time which is O(d log d) where d is the distance between the two split points. This value might be as large as X, but in total over all extend operations done in the tree, the sum of these distances can never be more than 2X, so the total time we get from these moves of the split points is O(n + X log X). Coming back to the ignored max(0, ...) part from above, the easiest way to deal with that is to simply remove any negative values from the end of the set after having performed the operation above.

We also need to revisit the join operation and see that this can still be done with the augmented data structure with shifts. Implementationwise, instead of keeping a single ordered set with a split point, it is probably easier (at least all our implementations seemed to feel so) to keep two separate sets corresponding to the part of values below and the part of values above the split point. Another implementation detail is that using complete ordered sets is not necessary, one can get away with two simple max-heaps.