Official

G - Frequency Query on Tree Editorial by Nyaan


この問題は 木上の Mo’s algorithm と呼ばれるテクニックで解くことが出来る問題です。

木上の Mo’s algorithm およびパスクエリの処理方法

この解説は Mo’s algorithm の知識を前提とします。Mo’s algorithm に関しては ABC293-G 解説 などを参照してください。

オイラーツアーをして、各頂点を行きがけと帰りがけに \(1\) 回ずつ記録することで出来る長さ \(2N\) の配列 \(E\) を考えます。また、各頂点が行きがけ・帰りがけのどのタイミングで記録されたかを示す配列 \(t_{in}, t_{out}\) も同時に記録します。具体的には (0-indexed/1-indexed の違いを無視すると)以下の Python 風擬似コードで記録される数列です。

# g : 木の隣接リスト表現
E = []
tin, tout = [-1] * N, [-1] * N

def dfs(c, p): 
  tin[c] = len(E)
  E.append(c)
  for d in g[c]:
    if d == p: continue
    dfs(d, c)
  tout[c] = len(E)
  E.append(c)

また、説明のためいくつかの記法を導入します。

  • 数列 \(A\) に対して \(A[L, R]\) を \(A\) の \(L\) 番目から \(R\) 番目までからなる数列とします。\(L \gt R\) の時 \(A[L,R]\) は空の数列です。
  • 数列 \(A\) に対して \(\mathrm{Odd}(A)\) を数列 \(A\) の中で奇数回登場する要素からなる整数集合とします。
  • 集合 \(S,T\) に対して \(S \triangle T\) を \(S\) と \(T\) の対称差、すなわち \(S\) と \(T\) の丁度一方だけに登場する要素からなる集合とします。

この時、以下の事実が成り立ちます。

頂点 \(S, T\) が \(t_{in}[S] \leq t_{in}[T]\) を満たす頂点対とする。この時、\(S, T\) パス上の頂点からなる頂点集合は、

\[\mathrm{Odd}(E[t_{in}[S] + 1, t_{in}[T]])\]

に \(\mathrm{LCA}(S,T)\) を追加したものと表すことが出来る。(ここで追加前の集合は \(\mathrm{LCA}(S,T)\) を含まないことが証明できる)

(証明)

\(S, T\) パス上の頂点からなる頂点集合を \(P\) とします。
まず、任意の頂点 \(S, T\) に対して

  • 根から頂点 \(S\) までに含まれる頂点からなる頂点集合を \(A\)
  • 根から頂点 \(T\) までに含まれる頂点からなる頂点集合を \(B\)

としてその対称差 \(A \triangle B\) を考えると、

\[A \triangle B = P \setminus \lbrace \mathrm{LCA}(S,T) \rbrace\]

となることが証明できます。(詳細略。根以外の各頂点に「その頂点を子とする辺」が対応することを利用して証明できます)

また、オイラーツアーの性質から \(A, B\) は \(E\) を用いて

\[A = \mathrm{Odd}(E[1, t_{in}[S]]), B = \mathrm{Odd}(E[1, t_{in}[T]])\]

となります。よって \(t_{in}[S] \leq t_{in}[T]\) の元で

\[A \triangle B = \mathrm{Odd}(E[t_{in}[S] + 1, t_{in}[T]])\]

が成り立ちます。(証明終わり)

この事実を利用すれば、各クエリを \(S, T\) から区間 \([t_{in}[S] + 1, t_{in}[T]]\) に変換することで木上のパスに対する Mo’s algorithm を適用できる形になります。具体的には以下の手順を行うことになります。

  • 各頂点 \(v\) が現在見ている範囲に登場する回数の偶奇を変数 parity[v] で管理する
  • 区間を伸縮するたびに parity[v] を更新して、その偶奇に応じて頂点 \(v\) に関する情報を追加あるいは削除する
  • クエリを処理する際は処理前に頂点 \(\mathrm{LCA}(S,T)\) に関する情報を追加して、処理後にそれを削除する。なお、\(\mathrm{LCA}\) は適切なデータ構造を別途用いて高速に計算することにする

今回の問題の解法

以上の内容を踏まえた上で今回の問題の解法を説明します。

まず、木上の Mo’s algorithm を適用することで今回の問題は以下の問題に帰着します。

各色の頻度列 \(f_1, f_2, \dots, f_N\) がある。はじめ全ての要素は \(0\) である。以下のクエリを高速に処理せよ。

  • \(f_i\) の値を \(1\) 増減させる。\(f_i\) は操作後も \(0 \leq f_i \leq N\) が成り立つ。
  • \(a,b\) が与えられるので \(a \leq f_i \leq b\) が成り立つ \(f_i\) の個数を出力する。

これは \(g_i = 「i 以上である f_j の個数」\) と置いて \(g_0, g_1, \dots, g_{N+1}\) を同時に管理することで \(\mathrm{O}(1)\) でクエリを処理できます。すなわち、以下の要領でクエリを高速に処理できます。

  • \(1\) 種類目のクエリによって変更される \(g_i\) の個数は \(1\) 個だけであるからその値を更新すればよい。
  • \(2\) 種類目のクエリの答えは \(g_a - g_{b+1}\) なのでその値を出力すればよい。

以上の内容を実装すれば今回の問題を \(\mathrm{O}(N \sqrt{Q} + Q \log N)\) 程度で計算できて、これは十分高速です。