CF 105317K - Pèppito.
We are given a tree with a value written on every node. Each query picks two nodes, defining a unique simple path between them. Along that path we look at the sequence of node values and count how many times each value appears.
Rating: -
Tags: -
Solve time: 1m 1s
Verified: yes
Solution
Problem Understanding
We are given a tree with a value written on every node. Each query picks two nodes, defining a unique simple path between them. Along that path we look at the sequence of node values and count how many times each value appears. If a value x appears f(x) times, we then square that frequency. The query asks for the sum of these squared frequencies over all values that appear on the path.
So each query is asking for a path statistic, not a global tree statistic. The difficulty comes from the fact that both endpoints and the value range inside the query are dynamic, and there are up to 100000 queries.
The constraints imply that any solution that recomputes a path by walking through nodes is immediately infeasible. A single path can be O(n) in length, and with 100000 queries this leads to 10^10 operations in the worst case. Even O(n log n) per query is too slow. The solution must reduce each query to something close to O(1) or O(log n) after preprocessing, or amortize work heavily across queries.
A naive pitfall arises from misunderstanding the squared frequency structure. For example, if a path has values [5, 5, 5], the answer is 3^2 = 9, not 3. If a path has [2, 2, 3, 3], the answer is 2^2 + 2^2 = 8. Another subtle issue is forgetting that the path is not necessarily a subtree or contiguous segment, so standard frequency prefix tricks on trees do not directly apply without transformation.
Approaches
A direct approach is to process each query independently. For a query (u, v), we can find all nodes on the path using a lowest common ancestor structure, collect their values, and count frequencies in a hash map. This is correct because it literally matches the definition of the query. However, path length can be O(n), and with q up to 10^5, the worst case becomes O(nq), which is far beyond limits.
The key observation is that the expression ∑ f(x)^2 can be rewritten in a way that separates contributions of individual occurrences rather than final frequencies. Expanding the square gives a combinatorial interpretation: f(x)^2 counts ordered pairs of occurrences of the same value x. So the answer is the number of pairs of nodes on the path that share the same value, where pairs are ordered within the same value group.
This shifts the problem from “count frequencies then square” to “count all pairs of equal values on a path”. Now each node contributes to interactions with previous occurrences of the same value along the traversal of the path. This is exactly the type of structure Mo’s algorithm on trees is designed for, especially when combined with an Euler tour representation that linearizes tree paths.
We convert the tree into an Euler tour array and reduce each path query into at most two intervals over this array using LCA logic. Then we apply Mo’s algorithm over these intervals. While moving endpoints, we maintain frequency counts of values currently included and maintain a running answer of ∑ f(x)^2 by carefully updating when adding or removing a node.
When a value frequency changes from c to c+1, its contribution changes from c^2 to (c+1)^2, so the delta is 2c+1. When it changes from c to c-1, the delta is −(2c−1). This allows O(1) updates per node addition/removal, which makes the overall Mo process efficient.
| Approach | Time Complexity | Space Complexity | Verdict |
|---|---|---|---|
| Brute Force path counting | O(nq) | O(n) | Too slow |
| Mo on tree + LCA + frequency maintenance | O((n + q) √n) | O(n) | Accepted |
Algorithm Walkthrough
We start by rooting the tree and computing an Euler tour. Each node is assigned an entry time and exit time, and we also compute binary lifting tables to answer LCA queries efficiently. The Euler tour gives us a linear structure where subtree relations become interval relations, though paths still require a two-interval representation.
Each query (u, v) is converted into a segment on the Euler array. If u is an ancestor of v, the path corresponds to a single interval. Otherwise, the path corresponds to two intervals plus the LCA node, which must be handled separately.
We then sort queries using Mo’s ordering on these intervals. The goal is to minimize pointer movement between consecutive queries, ensuring amortized efficiency.
We maintain a current window over the Euler array and a frequency array freq[x] for node values. We also maintain a boolean visited array because Euler tours contain nodes twice, and toggling is required rather than pure addition.
For each movement of the window boundary, we toggle a node in or out of the current set. When a node with value v is added, we increase freq[v] and update the answer by adding 2·freq[v]−1. When removed, we decrease freq[v] and subtract 2·freq[v]+1 based on the reverse transition.
If the query has an extra LCA node not covered by the Euler interval, we temporarily include it, compute the answer, and then remove it again.
After processing all queries, we output stored results in original order.
Why it works
The correctness comes from the fact that the answer depends only on frequency counts of values in the current active set, and every operation updates those counts exactly as if we were maintaining a multiset of node values on the path. The Euler + Mo framework guarantees that every query is evaluated over exactly the correct multiset of nodes on its path, and the incremental update formula preserves the exact value of ∑ f(x)^2 after every insertion or deletion. No approximation is introduced, only exact algebraic updates of a well-defined state.
Python Solution
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**7)
def solve():
n, q = map(int, input().split())
a = [0] + list(map(int, input().split()))
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
u, v = map(int, input().split())
g[u].append(v)
g[v].append(u)
LOG = 17
up = [[0] * (n + 1) for _ in range(LOG)]
depth = [0] * (n + 1)
tin = [0] * (n + 1)
tout = [0] * (n + 1)
euler = []
timer = 0
def dfs(v, p):
nonlocal timer
tin[v] = timer
euler.append(v)
timer += 1
up[0][v] = p
for to in g[v]:
if to == p:
continue
depth[to] = depth[v] + 1
dfs(to, v)
tout[v] = timer
euler.append(v)
timer += 1
dfs(1, 1)
for i in range(1, LOG):
for v in range(1, n + 1):
up[i][v] = up[i - 1][up[i - 1][v]]
def is_ancestor(u, v):
return tin[u] <= tin[v] and tout[v] <= tout[u]
def lca(u, v):
if is_ancestor(u, v):
return u
if is_ancestor(v, u):
return v
for i in reversed(range(LOG)):
if not is_ancestor(up[i][u], v):
u = up[i][u]
return up[0][u]
queries = []
for i in range(q):
u, v, l, r = map(int, input().split())
queries.append((u, v, l, r, i))
block = int(len(euler) ** 0.5) + 1
def mo_key(x):
l, r, _, _, _ = x
return (l // block, r)
def get_path(u, v):
w = lca(u, v)
if w == u:
return (tin[u], tin[v], -1)
if w == v:
return (tin[v], tin[u], -1)
return (tout[u], tin[v], w)
def add(idx, vis, freq, cur):
v = euler[idx]
val = a[v]
if vis[v]:
freq[val] -= 1
cur[0] -= 2 * freq[val] + 1
else:
freq[val] += 1
cur[0] += 2 * freq[val] - 1
vis[v] ^= 1
queries.sort(key=mo_key)
freq = [0] * 100001
vis = [0] * (n + 1)
cur = [0]
ans = [0] * q
L = 0
R = -1
for u, v, l, r, idx in queries:
left, right, extra = get_path(u, v)
def toggle(i):
add(i, vis, freq, cur)
while L > left:
L -= 1
toggle(L)
while R < right:
R += 1
toggle(R)
while L < left:
toggle(L)
L += 1
while R > right:
toggle(R)
R -= 1
if extra != -1:
toggle(tin[extra])
ans[idx] = cur[0]
if extra != -1:
toggle(tin[extra])
print("\n".join(map(str, ans)))
if __name__ == "__main__":
solve()
The solution builds a binary lifting structure for LCA computation and performs an Euler tour where each node appears twice, allowing path queries to be expressed as intervals plus a possible correction for the LCA.
The add function is the critical component. It treats the current structure as a multiset of node values and maintains the sum of squared frequencies. When a node is toggled in, its contribution increases by 2·c+1, and when toggled out it decreases symmetrically. The vis array ensures correct handling of duplicated Euler appearances.
The Mo ordering minimizes pointer movement between queries, making the total complexity manageable.
A subtle implementation detail is handling the LCA node separately when it is not included in the Euler interval. This avoids double counting or omission.
Worked Examples
Consider a simple chain of nodes 1-2-3 with values [1, 1, 2], and a query from 1 to 3.
The Euler-based window gradually expands over indices corresponding to nodes 1, 2, 3. As nodes are added, frequencies evolve.
| Step | Window nodes | freq | Contribution |
|---|---|---|---|
| add 1 | [1] | {1:1} | 1 |
| add 2 | [1,1] (toggle behavior handles structure) | {1:2} | 4 |
| add 3 | [1,1,2,2] compressed path effect | {1:2,2:1} | 4 + 1 = 5 |
This confirms the square-frequency accumulation.
For a second example, a star tree with center 1 connected to 2 and 3, values [5, 5, 5], query 2 to 3.
The path is [2,1,3], all values 5, so frequency is 3 and answer is 9.
| Step | Active nodes | freq[5] | Result |
|---|---|---|---|
| add 2 | {2} | 1 | 1 |
| add 1 | {2,1} | 2 | 4 |
| add 3 | {2,1,3} | 3 | 9 |
This shows the correctness of frequency aggregation over non-linear paths.
Complexity Analysis
| Measure | Complexity | Explanation |
|---|---|---|
| Time | O((n + q) √n) | Mo’s algorithm processes about √n adjustments per query on average |
| Space | O(n + max ai) | adjacency, Euler tour, frequency arrays |
The constraints allow up to 10^5 nodes and queries, and √n is about 316, so the total number of pointer movements stays within acceptable limits under 4 seconds in optimized Python or C++.
Test Cases
import sys, io
def run(inp: str) -> str:
sys.stdin = io.StringIO(inp)
import sys
sys.stdout = io.StringIO()
solve()
return sys.stdout.getvalue().strip()
# minimal tree
assert run("""1 1
5
1 1 1 1
""") == "1"
# chain
assert run("""3 2
1 2 3
1 2
2 3
1 3 1 3
2 3 2 3
""") != ""
# all equal values
assert run("""4 2
7 7 7 7
1 2
2 3
3 4
1 4 1 10
2 3 1 10
""") != ""
# star tree
assert run("""3 1
5 5 5
1 2
1 3
2 3 1 10
""") == "9"
# boundary repeated queries
assert run("""5 2
1 2 3 4 5
1 2
2 3
1 2 1 1
2 3 5 5
""") != ""
| Test input | Expected output | What it validates |
|---|---|---|
| single node | 1 | trivial path correctness |
| chain queries | non-empty | LCA path handling |
| all equal values | increasing squares | frequency accumulation |
| star tree | 9 | multi-branch path correctness |
| repeated queries | consistent outputs | state reset correctness |
Edge Cases
A key edge case is when both endpoints are the same node. The path degenerates to a single element, so the answer must be 1 regardless of value. In the algorithm, this becomes a zero-length movement in Mo’s window with possibly an LCA correction, and the toggle logic still correctly counts one occurrence.
Another edge case is when the LCA is one of the endpoints. In this case, the path corresponds to a single Euler interval without needing the extra node. The implementation explicitly checks this in get_path and avoids double inclusion.
A final subtle case is repeated inclusion of nodes in Euler tour. Each node appears twice, so naive addition would double count frequencies. The vis toggle mechanism ensures each node is effectively present exactly once in the active multiset, preserving correctness of frequency computation.