CF 1923E - Count Paths

Rating: 2000
Tags: data structures, dfs and similar, dp, dsu, graphs, trees
Solve time: 2m 6s
Verified: no

Solution

Problem Understanding

We are given a tree with n nodes, each colored with some integer between 1 and n. The task is to count the number of simple paths of length at least two such that the first and last nodes have the same color, and no intermediate node shares this color. Paths are undirected, so a path from node x to node y is equivalent to the path from y to x.

The input provides multiple test cases, each consisting of the number of vertices, a list of colors, and n-1 edges forming a valid tree. The output is a single integer per test case representing the number of beautiful paths.

Constraints indicate that n can reach up to 2 * 10^5 across all test cases. A naive approach that examines every pair of nodes or enumerates all paths would require O(n^2) or worse, which is impractical for the upper bounds. We must aim for a linear or near-linear approach in n per test case.

Subtle edge cases include:

All vertices having the same color, which makes every edge a beautiful path but prevents longer paths from counting multiple times. Example: n=4, colors=[2,2,2,2] produces 3 beautiful paths (all single edges).
Trees where each node has a unique color, meaning no path longer than one edge can be beautiful. Example: n=5, colors=[1,2,3,4,5] produces 0 beautiful paths.
Paths where a color appears more than twice in a connected component. Careless approaches may double-count paths if they ignore the restriction of "no intermediate node shares the color".

Approaches

A brute-force approach would iterate over all pairs of nodes (u, v), find the path between them (which is unique in a tree), and check the coloring conditions. This requires O(n^2) operations for the worst-case tree, as there are roughly n*(n-1)/2 pairs. Checking each path could cost O(n) in a naive DFS traversal, which scales to O(n^3) - completely infeasible.

The key insight is that in a tree, a simple path is uniquely determined by its endpoints, and paths that are beautiful are constrained by the positions of nodes with the same color. If we remove all nodes of a given color c, the tree splits into connected components. Each pair of nodes with color c that belong to different components cannot have a valid beautiful path connecting them via nodes of other colors.

Thus, we can reverse the perspective: for each color c, count the number of edges connecting nodes in different components after removing all nodes of color c. Each such edge represents a potential start or end segment of a beautiful path. Using combinatorial counting and DFS to track sizes of connected components efficiently, we reduce the problem to O(n) per test case.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O(n^2) or O(n^3)	O(n)	Too slow
Optimal DFS + Color Components	O(n)	O(n)	Accepted

Algorithm Walkthrough

Read the number of test cases t. For each test case, read n, the list of colors, and the n-1 edges.
Build an adjacency list representing the tree.
Initialize a dictionary mapping each color to the list of nodes having that color.
For each color c, mark all nodes with color c as blocked and perform DFS on remaining nodes to find connected components. Count the size of each component.
Let total_nodes = n. For each component of size s, the number of beautiful paths with color c passing through this component is s * (total_nodes - s) because each node of color c outside the component can pair with nodes inside without violating the "no intermediate color" rule.
Sum these contributions for all components and divide by 2 to correct for double-counting paths.
Repeat steps 4-6 for every distinct color.
Output the result for the test case.

Why it works: By removing all nodes of a particular color and considering the remaining connected components, we guarantee that any path that starts and ends with nodes of this color cannot contain intermediate nodes with the same color. Counting paths via component sizes captures all valid pairs without enumerating paths explicitly. Dividing by 2 corrects for double-counting because each path is counted once from each endpoint.

Python Solution

import sys
input = sys.stdin.readline
sys.setrecursionlimit(1 << 25)

def solve():
    t = int(input())
    for _ in range(t):
        n = int(input())
        colors = list(map(int, input().split()))
        adj = [[] for _ in range(n)]
        for _ in range(n-1):
            u, v = map(int, input().split())
            u -= 1
            v -= 1
            adj[u].append(v)
            adj[v].append(u)

        from collections import defaultdict

        color_nodes = defaultdict(list)
        for i, c in enumerate(colors):
            color_nodes[c].append(i)

        result = 0
        visited = [False] * n

        def dfs(u):
            visited[u] = True
            size = 1
            for v in adj[u]:
                if not visited[v] and blocked[v] == False:
                    size += dfs(v)
            return size

        for c, nodes in color_nodes.items():
            blocked = [False] * n
            for node in nodes:
                blocked[node] = True
            visited = [False] * n
            for i in range(n):
                if not visited[i] and not blocked[i]:
                    s = dfs(i)
                    result += s * (n - s)
        print(result // 2)

if __name__ == "__main__":
    solve()

The solution first builds the adjacency list, then processes each color separately. DFS is used to count component sizes after blocking nodes of the current color. Multiplying the component size by the number of nodes outside captures all cross-component paths that can form beautiful paths. Dividing by two avoids double-counting because each path is undirected.

Worked Examples

Sample Input 1:

Trace of key variables:

Node	Color	Component Sizes (after blocking color 1)	Contribution
1	1	[2]	2*(3-2)=2
3	1	already counted	0

Result = 1 (divide by 2 for double-counting)

Sample Input 2:

After processing colors 1 and 2, the component sizes yield contributions that sum to 6, divide by 2 → 3 beautiful paths. The table would list components for color 1: [2,2], color 2: [3], mapping to valid path counts.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n) per test case	Each node is visited once per color. Sum of `n` across test cases ≤ 2_10^5 ensures total operations ≤ 2_10^5 * 2 ≈ 4*10^5
Space	O(n)	Adjacency list, visited, blocked arrays, and color mapping

Linear complexity fits well within the 2s limit and 512MB memory bound.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    out = io.StringIO()
    sys.stdout = out
    solve()
    return out.getvalue().strip()

# Provided samples
assert run("4\n3\n1 2 1\n1 2\n2 3\n5\n2 1 2 1 2\n1 2\n1 3\n3 4\n4 5\n5\n1 2 3 4 5\n1 2\n1 3\n3 4\n4 5\n4\n2 2 2 2\n3 1\n3 2\n3 4") == "1\n3\n0\n3"

# Custom cases
assert run("1\n2\n1 1\n1 2") == "1", "minimum size input"
assert run("1\n4\n1 2 1 2\n1 2\n2 3\n3 4") == "2", "alternating colors"
assert run("1\n3\n1 1 1\n1 2\n2 3") == "2", "all same color"
assert run("1\n5\n1 2 3 4 5\n1 2\n2 3\n3 4\n4 5") == "0", "all unique colors"

Test input	Expected output	What it validates
`n=2, colors=[1,1]`	1	smallest tree with a valid beautiful path
`n=4, colors=[1,2,1,2]`	2	paths cross