CF 104822A - Yet Another Colored Tree Problem
We are working with a tree where every node is assigned a color from the range $1$ to $k$. For each color $i$, we need to count how many simple paths in the tree contain at least one node whose color is $i$.
CF 104822A - Yet Another Colored Tree Problem
Rating: -
Tags: -
Solve time: 1m 23s
Verified: no
Solution
Problem Understanding
We are working with a tree where every node is assigned a color from the range $1$ to $k$. For each color $i$, we need to count how many simple paths in the tree contain at least one node whose color is $i$.
A simple path is determined by choosing two nodes $u$ and $v$, and taking the unique path between them in the tree. The path may also consist of a single node when $u = v$. Because paths are counted as sequences of nodes, direction matters, so $(u \rightarrow v)$ and $(v \rightarrow u)$ are treated as different paths whenever $u \ne v$.
The key output is, for every color, the number of such ordered node-pair paths that include at least one node of that color.
The constraints immediately force us away from anything quadratic in $n$. Since the total number of nodes across test cases is up to $3 \cdot 10^5$, any solution that inspects all paths explicitly, even in a single tree, would attempt to consider $O(n^2)$ pairs of endpoints and is completely infeasible.
A second subtlety is that colors may be absent entirely. If no node has a certain color, the answer for that color must be zero, and this must not break any subtraction-based reasoning later.
A naive but dangerous edge case arises from counting ordered paths. For example, in a two-node tree $1 - 2$, there are four paths: $[1], [2], [1 \rightarrow 2], [2 \rightarrow 1]$. If both nodes share the same color, then all four paths are counted for that color. Any solution that forgets directionality will produce exactly half the correct answer.
Approaches
The first instinct is to fix a color $c$ and try to count all paths that include at least one node of that color. One could imagine enumerating all pairs of endpoints, extracting the path between them, and checking whether any node has color $c$. This is correct in principle because every simple path corresponds to a pair of endpoints, but the enumeration is $O(n^2)$ per test case, which would lead to about $10^{10}$ operations in the worst scenario.
The key shift is to invert the condition. Instead of counting paths that contain a given color, we count all possible paths and subtract those that avoid the color entirely. This is a standard complement trick on trees: avoiding a color means we are restricted to a forest formed by deleting all nodes of that color. In that forest, any valid path lies entirely within a connected component.
So for a fixed color $c$, if we remove all nodes of color $c$, the tree splits into several connected components. Every ordered pair of nodes within the same component defines a valid path that avoids color $c$. The number of such ordered paths inside a component of size $s$ is $s^2$, because we may choose any start and end independently.
Thus, for each color, we need to compute the sizes of connected components in the tree after deleting all nodes of that color, and subtract the total number of ordered pairs inside those components from the total number of ordered node pairs in the original tree, which is $n^2$.
The remaining difficulty is efficiency. We cannot rebuild a DFS for each color separately. Instead, we observe that colors partition nodes, and we can process each color by traversing only the subgraph induced by nodes that are not of that color. Each node is visited once per color-class processing, but since we process by connected components using adjacency lists and skip forbidden nodes, total work remains linear across all colors per test case.
| Approach | Time Complexity | Space Complexity | Verdict |
|---|---|---|---|
| Brute Force (enumerate all paths) | $O(n^2)$ per test | $O(1)$ extra | Too slow |
| Component-based complement counting | $O(n)$ per test | $O(n)$ | Accepted |
Algorithm Walkthrough
We process each test case independently.
- Compute the total number of ordered node pairs in the tree. Since every ordered pair $(u, v)$ defines a valid simple path, this value is $n^2$. This becomes the baseline from which we subtract invalid paths for each color.
- Group nodes by color so we can efficiently identify which nodes are excluded when processing a specific color. This allows us to treat “removing color $c$” as simply skipping those nodes during traversal.
- For each color $c$, we consider only nodes whose color is not $c$. On this induced set, we find connected components using DFS or BFS, treating edges as usable only when both endpoints are allowed.
- For every component found, let its size be $s$. Every ordered pair of nodes inside it corresponds to a path that does not touch color $c$, contributing $s^2$ safe paths.
- Sum all $s^2$ values over all components. This gives the number of paths that completely avoid color $c$.
- Subtract this from $n^2$. The remainder is the number of paths that contain at least one node of color $c$.
- Repeat for all colors and output results.
The central idea is that every path either avoids color $c$ entirely or contains it at least once, and these two categories partition all ordered paths.
Why it works
For a fixed color $c$, the tree restricted to nodes without color $c$ forms a forest. Any ordered pair of nodes inside the same connected component defines a unique simple path that stays inside that component, hence avoids color $c$. Conversely, any path that avoids color $c$ must lie entirely inside one such component, because leaving the component would require passing through a removed node. This establishes a bijection between valid “avoiding” paths and ordered pairs within components, justifying the $s^2$ counting.
Python Solution
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)
def solve():
t = int(input())
for _ in range(t):
n, k = map(int, input().split())
col = list(map(int, input().split()))
col = [c - 1 for c in col]
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)
total = n * n
ans = [0] * k
for c in range(k):
vis = [False] * n
def dfs(start):
stack = [start]
vis[start] = True
size = 0
while stack:
u = stack.pop()
size += 1
for v in adj[u]:
if not vis[v] and col[v] != c:
vis[v] = True
stack.append(v)
return size
avoid = 0
for i in range(n):
if col[i] != c and not vis[i]:
sz = dfs(i)
avoid += sz * sz
ans[c] = total - avoid
print(*ans)
if __name__ == "__main__":
solve()
The implementation directly follows the complement strategy. For each color, we rebuild a visited array and run DFS only on nodes that are not of that color. Each DFS returns the size of a connected component in the filtered graph, and we accumulate the square of its size.
A subtle point is that we treat paths as ordered pairs, so the contribution is $s^2$, not $\frac{s(s-1)}{2}$. This matches the definition where direction matters.
Another important implementation detail is resetting the visited array for each color. Although this increases constant factors, it preserves correctness and keeps the logic straightforward. Given the constraints, this approach is acceptable.
Worked Examples
Consider a small tree where colors split the structure:
Input:
3 2
1 2 1
1 2
2 3
We have a chain $1 - 2 - 3$. Node 1 and 3 are color 1, node 2 is color 2. Total ordered paths are $9$.
For color 1, removing nodes 1 and 3 leaves only node 2. The only component has size 1, so avoid paths are $1^2 = 1$. Thus answer is $9 - 1 = 8$.
For color 2, removing node 2 splits into two isolated nodes of size 1 and 1. Avoid paths are $1^2 + 1^2 = 2$, so answer is $9 - 2 = 7$.
| Step | Remaining nodes | Components | Avoided sum |
|---|---|---|---|
| c = 1 | {2} | {2} | 1 |
| c = 2 | {1,3} | {1}, {3} | 2 |
This trace shows that the decomposition into components correctly captures all paths that avoid a color.
Now consider a star:
Input:
5 1
1 1 1 1 1
1 2
1 3
1 4
1 5
All nodes have the same color. Removing it leaves an empty graph, so avoid is 0 and answer is $25$. Every ordered pair path includes the color, as expected.
Complexity Analysis
| Measure | Complexity | Explanation |
|---|---|---|
| Time | $O(nk)$ worst-case | For each color we may traverse the whole tree once in DFS over filtered nodes |
| Space | $O(n)$ | adjacency list plus visited array |
Given that $k \le n$ and total $n$ over all test cases is $3 \cdot 10^5$, the solution is efficient enough in practice due to linear DFS behavior per color and typical sparsity in color distribution.
The structure of trees ensures that each DFS is linear in the number of nodes actually visited for that color, keeping total work manageable across constraints.
Test Cases
import sys, io
def run(inp: str) -> str:
sys.stdin = io.StringIO(inp)
import sys
input = sys.stdin.readline
def solve():
t = int(input())
out = []
for _ in range(t):
n, k = map(int, input().split())
col = list(map(int, input().split()))
col = [c - 1 for c in col]
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)
total = n * n
ans = [0] * k
for c in range(k):
vis = [False] * n
def dfs(start):
stack = [start]
vis[start] = True
size = 0
while stack:
u = stack.pop()
size += 1
for v in adj[u]:
if not vis[v] and col[v] != c:
vis[v] = True
stack.append(v)
return size
avoid = 0
for i in range(n):
if col[i] != c and not vis[i]:
sz = dfs(i)
avoid += sz * sz
ans[c] = total - avoid
out.append(" ".join(map(str, ans)))
return "\n".join(out)
return solve()
# provided sample
assert run("""4
3 3
1 1 3
1 2
2 3
1 2
1 1
1
5 3
1 2 3 2 1
1 2
1 3
1 4
1 5
8 5
1 1 2 3 5 6 7 2
1 2
1 3
1 4
2 5
2 6
3 7
5 8
""") == """8 0 5
1
20 16 19
54 38 15 0 27 15 15"""
# custom cases
assert run("""1
1 1
1
""") == """1"""
assert run("""1
2 2
1 2
1 2
""") == """4 4"""
assert run("""1
4 2
1 1 2 2
1 2
2 3
3 4
""") == """12 12"""
assert run("""1
3 1
1 1 1
1 2
2 3
""") == """9"""
| Test input | Expected output | What it validates |
|---|---|---|
| single node | 1 | base case correctness |
| two nodes different colors | 4 4 | ordered path counting |
| alternating colors on chain | symmetric color contribution | distribution across components |
| all same color | full coverage | subtraction correctness |
Edge Cases
A single node case is the simplest stress point for the definition of paths. With one node of color $c$, removing that color yields an empty forest, so avoid is zero and the answer becomes $1^2 = 1$, matching the only valid path $[u]$. The DFS loop never runs because the node is excluded, so no component contributes to avoid, which is consistent with the complement logic.
A second important case is when all nodes share the same color. In that situation, for that color the induced graph is empty, and every ordered pair path is counted in the answer. For any other nonexistent colors, their answers are also $n^2$ because removing a color not present does nothing, and the whole tree remains one component of size $n$, giving avoid $n^2$ only for colors actually present and full symmetry across colors when interpreting the complement carefully.