Author:Saurabh Hota
Testor:Rajat De
Editorialist:Rishav Pati
Difficulty: Easy
Topics: Depth First Search, Counting
Problem:-
Given a tree of size $n$ ($1 \le n \le 10^5$) with each vertex colored either black or white. Find the number of paths with alternating colors.
Quick Explanation:-
Considering the values defined in the Explantion that follows, the number of paths in a subtree rooted at a black node numbered $n$ is:
$$ \sum f(c) + n_w + \sum w_c + \frac{1}{2} \sum (w_c \times (sw_n - w_c)) $$
Explanation:-
First let us root the entire tree at a random node, say node $1$.
Now suppose you have the value $f(n)$ which is the number of paths of alternating colors that
lie completely in the subtree of node number $n$.
Then we are asked to find the value of $f(1)$ in the question.
Now let us define some more terms:
$w_n$ : Number of alternating paths in the subtree of $n$ ending at $n$ where color of $n$ is white.
$sw_n$ : Sum of all $w_c$ where $c$ is a child of $n$.
$b_n$ : Number of alternating paths in the subtree of $n$ ending at $n$ where color of $n$ is black.
$sb_n$ : Sum of all $w_c$ where $c$ is a child of $n$.
$n_w$ : Number of white children of $n$.
$n_b$ : Number of black children of $n$.
Now, if we have all these values for every child of the node $n$ then we can find the value of $f(n)$ very easily.
Any path in $n$’s subtree is either completely inside a subtree of one of its children, or ends at node $n$ and all other nodes are inside the subtree of a child of $n$ or connects $2$ subtrees of $2$ different children of $n$.
All the paths that fall in the first category can be counted as: $\sum f(c)$ where $c$ varies over all children of $n$.
To count the number of paths that fall into the second category, first assume that $n$ is
black.
Now this quantity is simply $n_w + \sum w_c$.
$n_w$ new paths of length $2$ each and $\sum w_c$ by extending each path ending at a white child of
$n$ to the node $n$.
Also notice that each of these paths contribute exactly $1$ to the value of $b_n$, which gives us
the way to maintain $b_n$ and $w_n$ (which is $0$ in this case).
Now for all the paths that contain $n$ and lie in $2$ subtrees.
Since $n$ is black, both the adjacent nodes of $n$ in any such path must be white.
Hence, any such path is actually a concatenation of a white ending path of one child, $n$ and
a white ending path of another child.
So, this number is simply: $\frac{1}{2} \sum (w_c \times (sw_n - w_c))$.
Adding all three quantities gives us the value of $f(n)$ and the discussion also shows how to maintain all the required auxiliary values. The entire procedure can be implemented using a single Depth First Search from node number $1$. Hence, the time and memory complexity both are $O(n)$.