Heavy-Light Decomposition, spanning trees, segment tree
You’re given a tree on N vertices’s. Each edge of the tree is either black or white. A single-color subtree of this tree is a non-empty maximal set S of its edges such that all edges in S have the same color and each of these edges can be reached from any other edge in S by only traversing edges from S. A maximal set is a set to which we cannot add any edge without breaking this property.
You have to process Q queries. In each query you’re given two vertices u and v, you have to flip the color of all edges on the path from u to v and compute the number of single-color subtrees of the tree.
From subtrees to vertices’s
First consider the following problem. You’re given a graph with V vertices’s, you know that it’s spanning forest contains E edges. How many connected components does this graph have? Since each tree in spanning forest contains V_k-1 edges where V_k is the number of vertices’s in the tree, total number of edges in spanning forest is equal to V-C where C is the number of connected components in the graph. So total number of components is equal to C=V-E.
Now to our problem. Let’s consider each edge of our tree to be the vertex of some graph. There’s total of N-1 vertices’s in this graph. Now let’s implicitly build some kind of spanning forest on this graph. Consider some particular vertex v. Assume it has W_v incident white edges and B_v incident black edges. By definition all W_v white and all B_v black edges should lie in the same single-color subtree each.
Thus we may connect both groups in a tree using \max(W_v-1,0) and \max(B_v-1,0) edges correspondingly (for example we may simply place them on a line. In this way we’ll obtain spanning forest on the graph of edges. Indeed, we will have same connected components and all connected components will still be a tree because initial graph was a tree and we shouldn’t be afraid of any cycles appearing.
One of possible ways to see that there won’t be cycles is to consider the case when all edges are colored black. Then we will have 2(N-1)-N=N-2 edges in our graph which will be the single connected component thus it will be a spanning tree on N-1 vertices’s. And any other graph we may obtain by the construction described above may be also obtained from such graph for same colored edges by deleting some of edges in this spanning tree.
Now we should note that if for vertex v both W_v and B_v are at least 1 then it will contribute W_v+B_v-2 = \deg v-2 to the number of edges in spanning forest and otherwise if W_v=0 or B_v=0 it will contribute \max(W_v,B_v)-1=\deg v - 1 to the number of edges. Thus the total number of edges in spanning forest will be equal to (\sum\limits_v \deg v)-N-\sharp_2 where \sharp_2 is the number of vertices’s which have both black and white incident edges. Given that \sum\limits_v \deg v for tree is equal to 2(N-1) we obtain that total number of edges in spanning forest is N-2-\sharp_2. Given that number of vertices’s in edge graph is N-1 we obtain that there’s total of (N-1)-(N-2-\sharp_2)=1+\sharp_2 connected components in that graph.
In this way we reduced the problem to the following one:
- Invert colors of all edges on a path.
- Calculate the number of vertices’s which have incident edges of single same color.
Next useful thing to note is that instead of one query to path from u to v we can apply two queries to paths from u to the root and from v to the root as the edges from lca(u,v) to the root will be inverted twice and first inversion will be negated by the second one.
Now heavy-light decomposition should come in handy. If you’re not familiar with it, please familiarize yourself here, here and here. Assume that we already decomposed our tree in heavy paths and every non-leaf vertex has exactly one heavy edge going from it to its largest subtree. We may split all vertices’s in three groups:
- Vertices having equal (black) colors among all light edges descending from them
- Vertices having equal (white) colors among all light edges descending from them
- Vertices having distinct colors among light edges descending from them
In each query there might be only at most O(\log n) light edges affected (namely those on which ascending heavy paths end) thus we may maintain these groups efficiently having at most O(\log n) transfers from one group to another. All such transfers might be dealt with manually. After that we’re only interested in vertices of first kind. Again, we may split them in three groups:
- Vertices having equal (black) color of ascending edge and descending heavy edge
- Vertices having equal (white) color of ascending edge and descending heavy edge
- Vertices having distinct colors of ascending edge and descending heavy edge
Again, there might be only O(\log n) transfers from third group to first two groups or vice versa, namely for the vertex v in which ascending way to the root starts and in the ends of light edges on this path. They all should be dealt with manually. For all the other affected vertices’s, their group is simply inverted from 1 to 2 and vice versa.
Our task after each query is to calculate the number of vertices’s which are either in first or second group by both classification simultaneously. Since group via first classification tend to change very slowly, we may maintain two segment trees over heavy-light decomposition for first and second groups from first classification. Note that by construction of HLD you may make it so that all heavy paths form consequent segments in a segment tree, thus queries on ascending paths are reduced to queries on segments. Here’s full list of queries you’ll have to maintain:
- “Activate” or “deactivate” some specific vertex
- Calculate number of “activated” vertices’s having value “1” on the segment
- Invert values of all vertices’s on the segment from “1” to “2” and vice versa
As you may see, queries is quite standard for segment tree and may be easily implemented in O(\log N) per query. Note though that there will be O(\log N) heavy paths in each of Q initial queries thus total complexity of the solution will be as much as O(Q \log^2 N).