Note that there are many solutions - just two of them are listed here.

For the intended solution, we will use the term 'path' to imply that it contains more than one node, as this does not affect the answer.

Subtask 1

Instead of counting the number of paths with at least three nodes of the same colour, we will count the number of paths with less than three nodes of the same colour, and subtract this from the total number of paths. Let's refer to a path with less than three nodes of the same colour as a 'bad path'.

The total number of paths in the graph is ~\binom{N}{2} = \frac{N \times (N-1)}{2}~, since a path is uniquely determined by its endpoints.

For the first subtask, the only bad paths are those with length ~2~, which correspond to the edges of the graph. Since the tree contains ~N-1~ edges, the number of paths with less than three nodes is ~N-1~.

Hence, the answer is simply ~\frac{N \times (N-1)}{2} - (N-1)~.

Time complexity: ~\mathcal{O}(1)~

Subtask 2

By the pigeonhole principle, any path containing at least five nodes must contain three nodes of the same colour. Therefore, the only bad paths which we have yet to consider are with three or four nodes.

Since the graph is linear, there are only ~\mathcal{O}(N)~ paths with three or four nodes. Therefore, we can simply brute force over all paths with three or four nodes to determine the total number of bad paths.

Time complexity: ~\mathcal{O}(N)~

Subtask 3

In a general tree graph, it is no longer guaranteed that there are only ~\mathcal{O}(N)~ paths with three or four nodes, so brute force will be too slow.

Instead, we can precompute ~black[x]~ and ~white[x]~, which denote the number of black nodes and the number of white nodes which are adjacent to node ~x~ respectively.

To determine the number of bad paths with three nodes, we can iterate over ~x~, the middle of the three nodes. If node ~x~ is black, we add ~black[x] \times white[x] + \binom{white[x]}{2}~ to the number of bad paths seen so far. Otherwise, we can add ~black[x] \times white[x] + \binom{black[x]}{2}~ to the number of bad paths seen so far.

To determine the number of bad paths with four nodes, we can iterate over the middle two nodes, corresponding to the edges of the graph. By similar casework, depending on the colours of the middle two nodes, we can determine the number of bad paths with these two nodes in the middle using our precomputed values.

Time complexity: ~\mathcal{O}(N)~

Alternative Solution

Another solution is to use dynamic programming. Let ~dp[i][j][k]~ represent the number of paths with ~j~ black nodes and ~k~ white nodes such that node ~i~ is the LCA (Lowest Common Ancestor) of each of those paths. Note that since we only care about paths with up to ~3~ nodes of a colour, we can treat any path with more than ~3~ nodes of the same colour as a path with exactly ~3~ nodes of that colour. Also, note that we will be treating singular nodes as 'paths' for this solution.

To calculate the DP values at a node, we iterate through all of its children. In order to avoid going through all pairs of children, we instead merge each path from a child with all the paths from children that have already been processed. This will also ensure that we do not count any path with the same LCA more than once. The implementation of this process will be left as an exercise for the reader.

Overall, the total number of valid paths is ~\sum_{i=1}^n (\sum_{j=0}^2 dp[i][j][3] + \sum_{k=0}^3 dp[i][3][k])~.

We will not be over-counting any path since each path has exactly one LCA.

Time complexity: ~\mathcal{O}(N)~