Day 135-Q2 | 木上のMo's Algorithm（パスクエリ）

問題

$N$頂点の木があり、頂点$v$には色$C_v$（$1\le C_v\le N$）が塗られている。木は頂点$1$を根とする。$Q$個のクエリ$(u,v)$それぞれについて、$u$から$v$への単純パス上に現れる相異なる色の個数を答えよ。

入力形式

N
C_1 C_2 ... C_N
a_1 b_1
...
a_{N-1} b_{N-1}
Q
u_1 v_1
...
u_Q v_Q

制約

$2 \le N \le 10^5$

$1 \le C_i \le N$

木として連結

$1 \le Q \le 10^5$

$u_i \ne v_i$

入出力例

入力例1

出力例1

2
3

概念図

ヒント（段階的開示）

ヒント1（方向性）

木上のパスクエリはオイラーツアーで部分木には帰着できるが、LCAをまたぐパスは単純な区間にならない。「入退場時刻」トリックでパスクエリを配列上の区間クエリに変換し、通常のMo's algorithmを適用できる。

ヒント2（アプローチ）

各頂点の入場時刻in[v]と退場時刻out[v]を記録し、長さ$2N$の配列eulerを作る。$u$が$v$の祖先なら区間$[\text{in}[u],\text{in}[v]]$、そうでなければ$[\text{out}[u],\text{in}[v]]$とし、後者では$\text{LCA}(u,v)$を別途1回加える。

ヒント3（誘導）

区間内で頂点が2回現れたらパス外（トグルで打ち消される）、1回ならパス上という不変量を維持する。

def toggle(v):
    c = color[v]
    if visited[v]:
        cnt[c] -= 1
        if cnt[c] == 0:
            distinct[0] -= 1
        visited[v] = False
    else:
        if cnt[c] == 0:
            distinct[0] += 1
        cnt[c] += 1
        visited[v] = True

模範解答 (Python)

import sys
from math import isqrt


def main():
    data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    color = [0] + [int(data[idx + i]) for i in range(N)]
    idx += N

    adj = [[] for _ in range(N + 1)]
    for _ in range(N - 1):
        a = int(data[idx]); idx += 1
        b = int(data[idx]); idx += 1
        adj[a].append(b)
        adj[b].append(a)

    LOG = max(1, N.bit_length())
    parent = [[0] * (N + 1) for _ in range(LOG)]
    depth = [0] * (N + 1)
    tin = [0] * (N + 1)
    tout = [0] * (N + 1)
    euler = [0] * (2 * N)

    timer = 0
    visited_dfs = [False] * (N + 1)
    visited_dfs[1] = True
    it_stack = [iter(adj[1])]
    node_stack = [1]
    tin[1] = timer
    euler[timer] = 1
    timer += 1
    while node_stack:
        u = node_stack[-1]
        found = False
        for w in it_stack[-1]:
            if not visited_dfs[w]:
                visited_dfs[w] = True
                parent[0][w] = u
                depth[w] = depth[u] + 1
                tin[w] = timer
                euler[timer] = w
                timer += 1
                node_stack.append(w)
                it_stack.append(iter(adj[w]))
                found = True
                break
        if not found:
            tout[u] = timer
            euler[timer] = u
            timer += 1
            node_stack.pop()
            it_stack.pop()

    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(data[idx]); idx += 1
    queries = []
    for qi in range(Q):
        u = int(data[idx]); idx += 1
        v = int(data[idx]); idx += 1
        if tin[u] > tin[v]:
            u, v = v, u
        w = lca(u, v)
        if w == u:
            l, r, extra = tin[u], tin[v], -1
        else:
            l, r, extra = tout[u], tin[v], w
        queries.append((l, r, extra, qi))

    M = len(euler)
    block = max(1, isqrt(M))

    def sort_key(q):
        l, r, extra, qi = q
        b = l // block
        return (b, r if b % 2 == 0 else -r)

    queries.sort(key=sort_key)

    cnt = [0] * (N + 1)
    visited = [False] * (N + 1)
    distinct = 0

    def toggle(v):
        nonlocal distinct
        c = color[v]
        if visited[v]:
            cnt[c] -= 1
            if cnt[c] == 0:
                distinct -= 1
            visited[v] = False
        else:
            if cnt[c] == 0:
                distinct += 1
            cnt[c] += 1
            visited[v] = True

    cur_l, cur_r = 0, -1
    ans = [0] * Q
    for l, r, extra, qi in queries:
        while cur_r < r:
            cur_r += 1
            toggle(euler[cur_r])
        while cur_l > l:
            cur_l -= 1
            toggle(euler[cur_l])
        while cur_r > r:
            toggle(euler[cur_r])
            cur_r -= 1
        while cur_l < l:
            toggle(euler[cur_l])
            cur_l += 1
        if extra != -1:
            toggle(extra)
        ans[qi] = distinct
        if extra != -1:
            toggle(extra)

    print("\n".join(map(str, ans)))


main()

Step-by-Step 解説

1入退場時刻によるパス→区間変換
各頂点を入場・退場の2回配列に記録。区間内で奇数回登場する頂点集合＝パス上の頂点集合、という性質を利用。

2LCAの扱い
$u$が$v$の祖先でない場合、区間$[\text{out}[u],\text{in}[v]]$にはLCAが含まれないため別途1回加算する。

3トグル方式の相異なる値カウント
色ごとのカウントが$0\to1$でdistinctを増やし、$1\to0$で減らす不変量を維持する。

4Mo's Algorithm本体
配列長$2N$に通常のMo's algorithmを適用。ブロックサイズ$\sqrt{2N}$、ジグザグソートで定数倍を抑える。$O((N+Q)\sqrt N)$。

5LCAの一時トグル
答えを記録した直後にLCAを再トグルして元の状態に戻し、次のクエリの遷移を壊さない。

よくあるミス

ミス	原因	正しい書き方
LCAを2回カウント	区間に既に含まれるケースを見落として別加算	祖先関係で区間の取り方を正しく分岐する
extraの一時トグルを戻し忘れる	答え記録後に取り消しトグルを書かない	記録直後に必ず同じ頂点を再トグルする
反復DFSでtoutの記録漏れ	帰りがけ処理を書き忘れる	「子がもう無い」時に退場時刻を記録する
ブロックサイズをN基準で決める	配列長が$2N$であることを忘れる	block=isqrt(2N)にする
色カウントの境界ミス	cnt>0を条件に毎回加算	0→1、1→0の遷移でのみdistinctを更新する

次のステップ

発展: 特定の色$k$の出現回数を求めるクエリに変えた場合、同じ枠組みで解けるか検討せよ。
発展2: 頂点色の更新が入る場合、時間軸を追加した3次元Mo's algorithmへ拡張せよ。
次回予告: 次回のテーマは今後のセッションで選定

自己評価

完全理解おおむね理解要復習

Day 135-Q2 — 木上のMo's Algorithm（オイラーツアー圧縮によるパスクエリ）

問題