Interesting problem on tree

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:

1. $\text{value}[u]=\text{value}[u']$ for some $u'\neq u$ on the path from $r_i$ to $u$. In this case, we have already counted $\text{value}[u]$ before so just ignore it.
2. If not, we should count the number of $b$ such that $\text{value}[u]$ doesn't appear in $\text{path}(r,b)$ and add it to $AR[a]$. This is exactly $(n-n_i)-(|S(\text{value}[u])|-|S_i(\text{value}[u])|)$ where $n_i$ is the number of nodes in the $i$-th branch.

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$.

1. If $\text{value}[u]=\text{value}[u']$ for some $u'\neq u$ on the path from $r$ to $u$, we have already counted $\text{value}[u]$ before so just ignore it.
2. If not, add $n_u$ to $|S(\text{value}[u])|$. Here $n_u$ means the size of subtree rooted at $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)
• I understood what you're trying to achieve but don't understand the order in which you will perform the required steps. Will there be two DFS's on the centroid tree after it's construction for the two required pre computations, also if AR[r]=∑v|S(v)| , then why do we require all the steps after it as every node of the original tree will be root in some step of the centroid tree? It will be helpful if you explain the step by step algorithm for your solution so that it's easier to understand. – LTim Feb 13 '17 at 15:29
• Also, Pseudo code would be helpful (assuming, we already have the constructed centroid tree) – LTim Feb 13 '17 at 15:51
• @LTim Oh, I assumed you already know the standard divide and conquer procedure using centroid decomposition. After each phase we will remove the centroid, get some independent trees, and recurse on these trees. I'll fix it tomorrow and make the flow clearer. – aaaaajack Feb 13 '17 at 16:27
• @LTim Some details added. I suggest you don't accept the answer until you feel everything is clear. – aaaaajack Feb 14 '17 at 4:09

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.