Counting triplets with red edges in each pair

Question

Given a tree having N vertices and N-1 edges where each edges is having one of either red(r) or black(b) color. I need to find how many triplets(a,b,c) of vertices are there, such that on the path from vertex a to b, vertex b to c and vertex c to a there is atleast one edge having red color.

It should be noted that (a,b,c), (b,a,c) and all such permutation will be considered as the same triplets.

EXAMPLE : Let N=5 and edges with colors are as follow :

1 2 b
2 3 r
3 4 r
4 5 b

Here answer will be 4.

EXPLANATION : (2,3,4) is one such triplet because on all paths i.e 2 to 3, 3 to 4 and 2 to 4 there is atleast one edge having read color. (2,3,5), (1,3,4) and (1,3,5) are such other triplets.

Here is my approach :

Count the triplets that don't have this property, and subtract from the total number of triplets.

If you have a path made entirely of black edges, any third node will create such a triplet, so there are N-2 such triplets.

Counting the all-black paths involves removing all red edges and measuring the resultant black sub-trees (O(N)). A black tree containing K+1 nodes contains K(K+1)/2 paths.

Once you have the number of such triplets, subtract from the number of all triplets (N(N-1)(N-2)/6), and you have your answer in O(N).

But the problem in this approch is that some triplets will be counted multiple times.So how to handle is the problem

Also N can be upto 10^5 so i want a pretty fast algorithm for it.Almost O(N) OR O(NLOGN) time..not more than it

Answer 1

Let's call triplets having your property good triplets, and those that don't bad. Let's count the bad triplets. Let $\{T_1, T_2, \dots, T_m\}$ be the maximal subtrees with all black edges, and suppose these trees have $t_1, t_2, \dots, t_m$ vertices, respectively. For each $T_i$, the number of bad triplets with all three vertices among $V(T_i)$ is $\binom{t_i}{3}$. The number of bad triplets having two vertices among $V(T_i)$, and one vertex among the remaining vertices is $\binom{t_i}{2}(N - t_i)$. All bad triplets are of one of these forms. So we have that the number of good triplets is:

$$\binom{N}{3} - \sum_{i=1}^{m} \left(\binom{t_i}{3} + \binom{t_i}{2}(N - t_i) \right)$$ $$=\frac{N(N-1)(N-2)}{6} - \sum_{i=1}^{m} \left(\frac{t_i(t_i-1)(t_i-2)}{6} + \frac{t_i(t_i-1)}{2}(N - t_i) \right)$$ $$=\frac{1}{6}\left(N(N-1)(N-2) - \sum_{i=1}^{m} t_i(t_i-1)(3N-2t_i-2)\right)$$

Answer 2

1. Using depth-first search, create an array containing the sizes of connected components connected only by black edges in $\Theta(n)$ time and $\Theta(n)$ space.

2. Using the array, count the total number of triplets as the sum of product of all triplets in the array.

   The last step can be done by dynamic programming in $\Theta(m)$ time and $\Theta(1)$ space (if the array of sizes is of size $m$).

   • Let $s[k]$ be the sum of all combinations of $k$ elements in an array. If an item $x$ is added to array, the new value of $s[k]'$ is $s[k] + s[k-1] * x$ for $k>0$ with $s[0]=1$ for convenience.
   • Start with $s[k]=0$ for $k>0$ and add each item from A.
   • The answer is $s[3]$ after all elements have been added.