概要

区間和の情報を、累積和同士の差に関する制約として扱います。
版が木状に分岐するため、版の木を DFS しながら、ロールバック可能な重み付き Union-Find で制約を管理します。

考察

区間 \([L,R]\) のスコアは

\[ (A_L + A_{L+1} + \cdots + A_R) \bmod K \]

です。

ここで累積和を

\[ P_i = (A_1 + A_2 + \cdots + A_i) \bmod K \]

とおくと、区間 \([L,R]\) のスコアは

\[ P_R - P_{L-1} \pmod K \]

になります。

つまり、主張

\[ \text{score}(L,R) = X \]

は、次のような差分制約に言い換えられます。

\[ P_R - P_{L-1} \equiv X \pmod K \]

これは「頂点 \(L-1\) と頂点 \(R\) の値の差が \(X\)」という制約です。

---

制約をグラフとして見ると、各頂点は \(0,1,\dots,N\) の累積和 \(P_i\) に対応します。

ある 2 頂点 \(u,v\) が同じ連結成分に属していれば、これまでの制約から

\[ P_v - P_u \pmod K \]

が一意に決まります。

逆に、異なる連結成分に属しているなら、一方の成分全体を自由にずらせるため、差は一意に決まりません。したがって UNKNOWN です。

---

素朴に各版ごとに制約集合や Union-Find をコピーすると、版の数が最大 \(Q\) 個あるため、\(O(NQ)\) などになってしまい間に合いません。

そこで重要なのが、操作によって作られる版を木として見ることです。

各操作 \(i\) は既存の版 \(B\) を参照して新しい版 \(i\) を作るので、版 \(B\) から版 \(i\) へ辺を張ると、版全体は根を版 \(0\) とする木になります。

この木を DFS しながら、

• 木を下るときに制約を追加する
• 木を戻るときに追加した制約を取り消す

という処理をすれば、各版に対応する状態を効率よく再現できます。

このために、ロールバック可能な重み付き Union-Find を使います。

アルゴリズム

まず、各操作を読み込み、参照先の版 \(B\) から現在の版 \(i\) へ辺を張って、版の木を作ります。

各区間 \([L,R]\) は、累積和の頂点

\[ a = L-1,\quad b = R \]

に対応します。

種類 \(0\) の操作は制約

\[ P_b - P_a \equiv X \pmod K \]

を追加できるか判定します。

種類 \(1\) の操作は

\[ P_b - P_a \pmod K \]

が一意に決まるかを調べます。

---

重み付き Union-Find では、各頂点について親への差分を持ちます。

find(v) では、根とともに

\[ P_v - P_{\mathrm{root}} \pmod K \]

を返します。

制約

\[ P_b - P_a \equiv X \pmod K \]

を追加する場合を考えます。

• \(a,b\) がすでに同じ連結成分にある場合
  すでに決まっている差分が \(X\) と一致すれば受理、違えば矛盾なので却下します。

• \(a,b\) が異なる連結成分にある場合
  2 つの成分の相対位置はまだ自由なので、必ず制約を追加できます。
  Union-Find で併合します。

---

版の木を DFS します。

版 \(v\) に入る直前の Union-Find の状態は、親版の状態です。

版 \(v\) の操作を処理します。

• 種類 \(0\) の場合
  • 制約を追加できれば YES
  • 矛盾するなら NO
• 種類 \(1\) の場合
  • 両端が同じ連結成分なら差分を出力
  • 異なるなら UNKNOWN

その後、子の版へ進みます。

DFS で版 \(v\) の部分木をすべて処理し終えたら、版 \(v\) に入る前の状態まで Union-Find をロールバックします。

---

例えば、制約 0 B L R X は、Union-Find 上では

\[ P_R - P_{L-1} \equiv X \pmod K \]

という辺を追加する操作になります。

クエリ 1 B L R は、頂点 \(L-1\) と \(R\) が同じ連結成分にあるかを調べ、同じならその差分を答えます。

計算量

• 時間計算量: \(O(N + Q \log N)\)
• 空間計算量: \(O(N + Q)\)

Union-Find はロールバックのために経路圧縮を使いませんが、サイズによる併合を行うため、高さは \(O(\log N)\) に抑えられます。

実装のポイント

ロールバック可能な Union-Find では、経路圧縮をしてはいけません。
経路圧縮をすると多くの親情報が書き換わり、元に戻すのが難しくなるためです。

その代わり、サイズが小さい木を大きい木に併合することで、高さを抑えます。

また、併合したときには履歴に

• どの根を子にしたか
• どの根を親にしたか
• 親側の併合前のサイズ

を保存しておきます。

DFS から戻るときに、この履歴を使って Union-Find を元の状態に戻します。

コードでは再帰 DFS ではなく、明示的なスタックを使っています。
これは、\(Q\) が最大 \(10^5\) と大きく、再帰だと Python の再帰上限に引っかかる可能性があるためです。

ソースコード

import sys

def main():
    input = sys.stdin.buffer.readline
    N, K, Q = map(int, input().split())

    children = [[] for _ in range(Q + 1)]
    typ = [0] * (Q + 1)
    aa = [0] * (Q + 1)
    bb = [0] * (Q + 1)
    xx = [0] * (Q + 1)

    for i in range(1, Q + 1):
        q = list(map(int, input().split()))
        t = q[0]
        B = q[1]
        typ[i] = t
        children[B].append(i)
        if t == 0:
            _, _, L, R, X = q
            aa[i] = L - 1
            bb[i] = R
            xx[i] = X
        else:
            _, _, L, R = q
            aa[i] = L - 1
            bb[i] = R

    parent = list(range(N + 1))
    size = [1] * (N + 1)
    weight = [0] * (N + 1)
    history = []

    def find(v):
        s = 0
        while parent[v] != v:
            s += weight[v]
            if s >= K:
                s -= K
            v = parent[v]
        return v, s

    def unite(a, b, x):
        ra, pa = find(a)
        rb, pb = find(b)

        if ra == rb:
            d = pb - pa
            if d < 0:
                d += K
            return d == x

        if size[ra] < size[rb]:
            history.append((ra, rb, size[rb]))
            parent[ra] = rb
            weight[ra] = (pb - pa - x) % K
            size[rb] += size[ra]
        else:
            history.append((rb, ra, size[ra]))
            parent[rb] = ra
            weight[rb] = (x + pa - pb) % K
            size[ra] += size[rb]

        return True

    def rollback(snap):
        while len(history) > snap:
            child, par, old_size = history.pop()
            parent[child] = child
            weight[child] = 0
            size[par] = old_size

    ans = [""] * (Q + 1)
    stack = [(0, 0, 0)]

    while stack:
        v, state, snap = stack.pop()

        if state == 1:
            rollback(snap)
            continue

        snap = len(history)

        if v:
            a = aa[v]
            b = bb[v]

            if typ[v] == 0:
                if unite(a, b, xx[v]):
                    ans[v] = "YES"
                else:
                    ans[v] = "NO"
            else:
                ra, pa = find(a)
                rb, pb = find(b)
                if ra == rb:
                    d = pb - pa
                    if d < 0:
                        d += K
                    ans[v] = str(d)
                else:
                    ans[v] = "UNKNOWN"

        stack.append((v, 1, snap))
        for c in reversed(children[v]):
            stack.append((c, 0, 0))

    print("\n".join(ans[1:]))

if __name__ == "__main__":
    main()

---

この解説は gpt-5.5-high によって生成されました。