問題
N頂点の根付き木が与えられ、各頂点に色 c[v] が付いている。Q個のクエリで2頂点 u, v が指定される。u から v へのパス上に存在する異なる色の数を答えよ。
制約
$N, Q \le 5 \times 10^4$
$1 \le c[v] \le N$
木は連結
入出力例
入力例 1
7
1 2 3 1 2 3 4
1 2
1 3
2 4
2 5
3 6
3 7
3
4 6
2 7
1 4出力例 1
3
4
2ヒント (段階的開示)
ヒント1: 方向性
木のパス上のクエリをMo法で処理。オイラーツアーで木を1次元に変換し、パス(u,v)をオイラーツアー上の区間として表現。
ヒント2: アプローチ
1. オイラーツアーで in/out 時刻を記録。2. u, v のパスを [in[u], in[v]] または [out[u], in[v]] の区間に変換（LCA考慮）。3. Mo法でこの区間を処理。4. ビジットフラグで「現在パスに属するか」を管理。
ヒント3: 誘導
# lca = LCA(u, v)
# if lca == u: query range = [in[u], in[v]]
# else: query range = [out[u], in[v]], include lca separately模範解答 (Python)
import sys
from collections import defaultdict, deque
input = sys.stdin.readline
def main():
N = int(input())
color = [0] + list(map(int, input().split()))
graph = defaultdict(list)
for _ in range(N-1):
u, v = map(int, input().split())
graph[u].append(v)
graph[v].append(u)
euler = []
in_time = [0] * (N+1)
out_time = [0] * (N+1)
depth_arr = [0] * (N+1)
timer = [0]
stack = [(1, 0, False)]
while stack:
v, p, returning = stack.pop()
if returning:
euler.append(v); out_time[v] = timer[0]; timer[0] += 1
else:
in_time[v] = timer[0]; euler.append(v); timer[0] += 1
stack.append((v, p, True))
for u in graph[v]:
if u != p:
depth_arr[u] = depth_arr[v] + 1
stack.append((u, v, False))
L = len(euler)
LOG = 18
parent = [[0]*(N+1) for _ in range(LOG)]
depth = depth_arr[:]
bfs = deque([1])
visited = [False]*(N+1); visited[1] = True
par = [0]*(N+1)
while bfs:
v = bfs.popleft()
parent[0][v] = par[v] if par[v] != 0 else v
for u in graph[v]:
if not visited[u]:
visited[u] = True; par[u] = v; bfs.append(u)
parent[0][1] = 1
for k in range(1, LOG):
for v in range(1, N+1):
parent[k][v] = parent[k-1][parent[k-1][v]]
def lca(u, v):
if depth[u] < depth[v]: u, v = v, u
diff = depth[u] - depth[v]
for k in range(LOG):
if (diff >> k) & 1: u = parent[k][u]
if u == v: return u
for k in range(LOG-1, -1, -1):
if parent[k][u] != parent[k][v]:
u = parent[k][u]; v = parent[k][v]
return parent[0][u]
Q = int(input())
queries = []
for i in range(Q):
u, v = map(int, input().split())
l = lca(u, v)
if in_time[u] > in_time[v]: u, v = v, u
if l == u:
queries.append((in_time[u], in_time[v], -1, i))
else:
queries.append((out_time[u], in_time[v], l, i))
BLOCK = max(1, int(L**0.5))
queries.sort(key=lambda x: (x[0]//BLOCK, x[1] if (x[0]//BLOCK)%2==0 else -x[1]))
on_path = [False] * (N+1)
cnt = [0] * (N+1)
distinct = [0]
def toggle(v):
if on_path[v]:
on_path[v] = False
cnt[color[v]] -= 1
if cnt[color[v]] == 0: distinct[0] -= 1
else:
on_path[v] = True
if cnt[color[v]] == 0: distinct[0] += 1
cnt[color[v]] += 1
ans = [0] * Q
cur_l, cur_r = 0, -1
for ql, qr, qlca, qi in queries:
while cur_r < qr: cur_r += 1; toggle(euler[cur_r])
while cur_l > ql: cur_l -= 1; toggle(euler[cur_l])
while cur_r > qr: toggle(euler[cur_r]); cur_r -= 1
while cur_l < ql: toggle(euler[cur_l]); cur_l += 1
if qlca != -1: toggle(qlca)
ans[qi] = distinct[0]
if qlca != -1: toggle(qlca)
print('\n'.join(map(str, ans)))
main()Step-by-Step 解説
1オイラーツアー変換
各頂点を「行き」と「帰り」の2回登録。長さ 2N-1 の配列。
各頂点を「行き」と「帰り」の2回登録。長さ 2N-1 の配列。
2パスの区間変換
l == u なら [in[u], in[v]]、l != u なら [out[u], in[v]] + LCA を別途追加。
l == u なら [in[u], in[v]]、l != u なら [out[u], in[v]] + LCA を別途追加。
3Mo法でのトグル
頂点をパスに追加/除去するときビジットフラグでオンオフを切り替える。
頂点をパスに追加/除去するときビジットフラグでオンオフを切り替える。
4異なる色の数の管理
色のカウント配列と distinct カウンタで O(1) の追加/削除。
色のカウント配列と distinct カウンタで O(1) の追加/削除。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| LCAをトグルし忘れ | l != u の場合にLCAがパスに含まれない | クエリ処理後にlcaを一時的にトグル |
| オイラーツアーの長さ | 2N-1 を確認 | 各内部頂点の子数分だけ2倍 |
| Mo法のブロックサイズ | √(2N-1) が適切 | BLOCK = int((2N)**0.5) |
次のステップ
- 発展: パス上の色数の代わりに最頻値を求める
- 応用: 木上のMo法 + セグメント木の組み合わせ