10672 - Marbles on a Tree
We are given a rooted tree with n vertices. For each vertex v, we are given m marbles which rest on v. In total, there are n marbles. Our goal is to equally distribute the n marbles, such that we minimize the number of marble movements needed to reach equilibrium.
The first thing to notice is that n may be as large as 10,000, and thus we will need a reasonably fast solution. An algorithm that takes Θ(n2) time will probably timeout. Initially, it is hard to come up with an optimal strategy for moving the marbles, it seems that there are so many possible choices. Should we try to move all the excess marbles from the vertex with the greatest marble count? Where would we put them? However, if we consider a leaf vertex v in the tree, it is clear that they only way to move marbles to and from v is to get them from the parent of v. Thus, if leaf vertex v has m marbles, clearly, we must move m-1 marbles from v to the parent of v, so that v will have exactly 1 marble.
This idea gives us a nice algorithm to solve the problem. While the tree has more than one node, we take a leaf, move all but one of its marbles to its parent, and remove the leaf. While it is possible that a vertex will have a negative marble count during the process of this algorithm that is perfectly acceptable, it just means that that vertex is 'owed' marbles by its parent so that it can satisify its children. The sum of the marble count through the algorithm always equals exactly n, so there are still enough marbles to go around.
We maintain the tree as a simple data structure as follows. We keep two arrays, One array stores at index i, the parent of vertex i. The other array stores at index i, the out degree of vertex i. Note that the leaves in the tree are vertices having out degree 0.
Now all that is needed is a way to iterate through the leaves of a rooted tree, removing the the leaves from the tree are visited. This is really just the reverse of a topological sort. We do something resembling a breadth first search from the leaves up to the root. We create a queue, initially empty, and then proceed to put all leaves in the queue. We find the leaves with a simple linear search through the out degree array. Then we begin dequeing leaves from our queue. When we take a new leaf v, from the front of the queue. We decrement its parent's out degree by 1, since v will be removed from the tree, its parent will have one less child. If the outdegree of the parent becomes 0, we put the parent in the back of the queue, as the parent just became a leaf. Now we do any needed processing on v. In our case, our leaf v has m marbles, so we move m-1 marbles to our parent, and we add | m - 1 | to our move count.
Consider an arbitrary edge of the tree. Removing it splits the tree into two subtrees.
Let there be A marbles in one of the subtrees, let B be the number of nodes in that subtree. Then clearly in an optimal solution exactly |B-A| marbles will travel along this edge. Summing this up for all edges gives the desired answer. This leads to an alternate solution.
This solution can also be implemented as a simple depth-first search of the tree. Once we are leaving a vertex v, we can compute the size and marble count for its subtree, and thus the count of marbles traveling along the edge from v to its parent.
9 1 2 3 2 3 4 2 1 0 3 0 2 5 6 4 1 3 7 8 9 5 3 0 6 0 0 7 0 0 8 2 0 9 0 0 9 1 0 3 2 3 4 2 0 0 3 0 2 5 6 4 9 3 7 8 9 5 0 0 6 0 0 7 0 0 8 0 0 9 0 0 9 1 0 3 2 3 4 2 9 0 3 0 2 5 6 4 0 3 7 8 9 5 0 0 6 0 0 7 0 0 8 0 0 9 0 0 0
7 14 20