You are given a tree of $N$ nodes with colors assigned to each node. You need to find for each node $X$, a value $AR[X]$.
$AR[X]$ = The sum of values of $CALC(X , i)$ for each $i$ from $1$ to $N$.
$CALC(A,B)$ = Count the number of non-repeating colors in the path from $A$ to $B$.
Can somebody suggest an efficient solution to this problem?
For a fixed root, there are two kinds of path: the ones which pass the root and the ones which don't. The first kind is special since we can seperate a path $(a,b)$ into $(a,r)$ and $(r,b)$, which might be easier to deal with via proper preprocessing. Consider following divide and conquer procedure:
Solve(T):
Fix a root r
for each a
for each b s.t. (a,b) pass r
AR[a]+=calc(a,b)
for each branch T_i of r (which is also a tree)
Solve(T_i)
Note that we neither miss a pair nor count a pair twice since $(a,b)$ don't pass $r$ if and only if they are in the same branch. In addition, it's not hard to prove that trees in depth-$d$ recursive call don't overlap. Therefore, if we can do the non-recursive part in $O(\text{size of tree})$, the complexity would be $O(\text{recursive depth}\times n)$. If we always choose $r$ as the centroid, recursive depth is bounded by $O(\log n)$ and hence the complexity would be $O(n \log n)$. Now let's see how to solve the non-recursive part.
For a pair of node $(a,b)$, let $\text{path}(a,b)$ denote the set of distinct values on the path from $a$ to $b$, and $I_{ab}(v)$ denote the indicator function which is $1$ when $v\in \text{path}(a,b)$ and $0$ otherwise. It's not hard to see $\text{calc}(a,b)=|\text{path}(a,b)|=\sum_vI_v$.
Recall that in each phase, we choose the centroid $r$ as the root and consider only the paths passing the root. Let $B_a=\{b\in T:(a,b)\text{ pass }r\}$. Using the indicator function above, we have $\sum_{b\in B_a}\text{calc}(a,b)=\sum_{b\in B_a}\sum_v I_{ab}(v)=\sum_v\sum_{b\in B_a} I_{ab}(v)$. Therefore, computing $\sum_{b\in B_a} I_{ab}(v)$ for each $a$ and $v$ also gives us the answer. (You can interpret this as the number of $b$ such that $v\in \text{path}(a,b)$.) To do this, consider the set $S$ for each value $v$ such that $S_T(v)=\{u\in T: v\in\text{path}(r,u)\}$. We need to maintain $|S(v)|$ for the whole tree and similarly $|S_i(v)|$ for the $i$-th branch of $r$. We will see how to implement this later.
Since $B_r=T$, $\sum_{b\in B_r}\text{calc}(r,b)$ is exactly $\sum_v|S(v)|$. Next, consider node $a$ in the $i$-th branch. Let $r_i$ denote the $i$-th child of $r$. For each node $b$ which is not in the $i$-th branch, we can split path $(a,b)$ into $(r,b)$ and $(r_i,a)$. Therefore, $$\sum_{b\in B_a}\text{calc}(a,b)=\sum_{b\in B_a}\text{calc}(r,b)+\sum_{b\in B_a}|\text{path}(r_i,a)\backslash\text{path}(r,b)|$$ It's not hard to see that the first term equals to $\sum_v(|S_T(v)|-|S_i(v)|)$, so let's focus on the second term. Now consider each node $u$ on the path from $r_i$ to $a$ and count $\text{value}[u]$ one by one. There are two possible cases:
Note that the number we count for each $u$ is independent from $a$. Therefore we can do this for every $a$ in a single DFS which takes $O(n)$. (Tell me if you need more details.)
Now the only remaining part is the preprocessing which computes $|S(v)|$, and it's actually similar to previous steps:
Consider each node $u$.
$|S_i(v)|$ can be computed similarly. The preprocessing also takes $O(n)$.
The complete flow seems like
Solve(T):
Choose a centroid r as the root
Preprocess() //Compute |S(v)| and |S_i(v)| via a tree traversal
AR[r]+=Sum{|S(v)|}
for each node u
AR[u]+=Sum{|S(v)|-|S_i(v)|}
Update AR with the second term via another tree traversal
for each branch T_i of r (which is also a tree)
Solve(T_i)
Let's call a number of unique values over path $(a, b)$ its cost. You need at first to calculate a matrix of costs for all possible pairs in the tree. This matrix will be symmetric with ones on the main diagonal. After you have this matrix, you can find your $AR(n)$ values by summarizing costs on its rows (or columns).
Let's at first build a matrix $M$ of sets - each element of this matrix, corresponding to a path $(a, b)$ will contain a set $S(a, b)$ of (unique) values on this path. The cost of the path $(a, b)$ will be a number of elements in the set $S(a, b)$.
The following recursive process will produce this matrix of sets:
Subdivbision step. Find an edge $(b_1, b_2)$ in the tree $T$ in such way that its removal will give you two trees $T_1$ and $T_2$ with approximately the same number of nodes. Calculate matrices of sets for both $T_1$ and $T_2$.
Merging step. The matrix of sets for the tree $T$ will consist of four parts - the matrix of sets for $T_1$, the matrix of sets for $T_2$, and two "cross-over" matrices for paths from $T_1$ to $T_2$ and for paths from $T_2$ to $T_1$ (you can calculate just one of the cross-over matrices because of symmetry). The following formula can be used to calculate a set $S(n_1, n_2)$:
$$S(n_1, n_2) = S(n_1, b_1) \cup S(b_1, b_2) \cup S(b_2, n_2)$$
where $n_1 \in T_1$ and $n_2 \in T_2$.
You can use Tree Centroid Decomposition during the subdivision step - and I'll leave that to you.
