Day 134-Q1 — 重み付きUnion-Find(法P上の相対値制約管理)

2026-08-26 赤色 Master / Phase 8+ ★★★★★★★★★ 巡回群Z/PZ上のポテンシャル付きDSU O(α(N))

問題

$N$個の変数 $x_1,\dots,x_N$(それぞれ $\mathbb{Z}/P\mathbb{Z}$ 上の値、初期状態では未確定)についての情報が$Q$個、順番に与えられる。次の2種類のクエリを処理せよ。

  • 1 a b w : 制約 $x_b - x_a \equiv w \pmod P$ を追加する。既存の制約と矛盾する場合は無視する(構造も変化させない)。矛盾しなければ採用する。
  • 2 a b : 現時点で $x_b - x_a \pmod P$ が一意に定まるならその値を、定まらなければ -1 を出力する。

入力形式

N P Q
type_1 ...
type_2 ...
...
type_Q ...

(type_iが1なら続けてa b w、2なら続けてa b)

制約

$2 \le N \le 2\times10^5$
$2 \le P \le 10^9+7$(素数とは限らない)
$1 \le Q \le 2\times10^5$
$1 \le a,b \le N$、$0 \le w < P$

入出力例

入力例1

4 5 8
1 1 2 3
1 2 3 1
2 1 3
1 1 3 0
2 1 3
1 4 1 2
2 4 3
2 2 4

出力例1

4
4
1
0

概念図

根への合成: pot[x] = 「x の値 − 根の値」(mod P) root x1 pot=3 x2 pot=4 x3 親経由(圧縮前) 経路圧縮で直接張替え 旧親x2 find(x3): 経路 x3→x2→root のpotを合算しmod Pで正規化、根に直接付け替え

ヒント(段階的開示)

ヒント1(方向性)

通常の重み付きUnion-Find(整数の差分を管理するタイプ)と違い、今回は差分が $\mathbb{Z}/P\mathbb{Z}$ 上、つまり「$P$で割った余り」の世界に閉じている。加減算はすべてmod $P$で行う必要があり、符号反転や経路合成のすべてでmodを取り忘れると負の値や$P$以上の値が紛れ込む。

ヒント2(アプローチ)

各頂点$x$に「自分の値$-$親の値」を表すポテンシャルpot[x](mod $P$)を持たせ、findで根までのポテンシャルを合成しながら経路圧縮する。2つの根ra, rbを併合する際は、既知のポテンシャルpa=pot[a]、pb=pot[b]と与えられた$w$($x_b-x_a=w$)から、根同士の相対値を代数的に導出する。

ヒント3(誘導)

$x_a-x_{r_a}=p_a$、$x_b-x_{r_b}=p_b$、$x_b-x_a=w$から

$$x_{r_b}-x_{r_a}=(x_b-p_b)-(x_a-p_a)=w+p_a-p_b \pmod P$$

再帰は使わずスタックで経路を記録してから一括圧縮する反復版を書くこと。

def find(x):
    path = []
    edge_pots = []
    cx = x
    while par[cx] != cx:
        path.append(cx)
        edge_pots.append(pot[cx])
        cx = par[cx]
    root = cx
    suffix = 0
    for i in range(len(path) - 1, -1, -1):
        suffix = (suffix + edge_pots[i]) % P
        par[path[i]] = root
        pot[path[i]] = suffix
    return root, (pot[x] if path else 0)

模範解答 (Python)

import sys


def main():
    data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(data[idx]); idx += 1
    P = int(data[idx]); idx += 1
    Q = int(data[idx]); idx += 1

    par = list(range(N))
    rnk = [0] * N
    pot = [0] * N  # pot[x] = (value[x] - value[parent[x]]) mod P

    def find(x):
        path = []
        edge_pots = []
        cx = x
        while par[cx] != cx:
            path.append(cx)
            edge_pots.append(pot[cx])
            cx = par[cx]
        root = cx
        suffix = 0
        for i in range(len(path) - 1, -1, -1):
            suffix = (suffix + edge_pots[i]) % P
            node = path[i]
            par[node] = root
            pot[node] = suffix
        p = pot[x] if path else 0
        return root, p

    def union(a, b, w):
        ra, pa = find(a)
        rb, pb = find(b)
        if ra == rb:
            return ((pb - pa) % P) == (w % P)
        if rnk[ra] < rnk[rb]:
            ra, rb = rb, ra
            pa, pb = pb, pa
            w = (-w) % P
        newpot = (w + pa - pb) % P
        par[rb] = ra
        pot[rb] = newpot
        if rnk[ra] == rnk[rb]:
            rnk[ra] += 1
        return True

    def diff(a, b):
        ra, pa = find(a)
        rb, pb = find(b)
        if ra != rb:
            return None
        return (pb - pa) % P

    out = []
    for _ in range(Q):
        t = int(data[idx]); idx += 1
        if t == 1:
            a = int(data[idx]) - 1; idx += 1
            b = int(data[idx]) - 1; idx += 1
            w = int(data[idx]); idx += 1
            union(a, b, w)
        else:
            a = int(data[idx]) - 1; idx += 1
            b = int(data[idx]) - 1; idx += 1
            r = diff(a, b)
            out.append(str(r) if r is not None else "-1")
    print("\n".join(out))


main()

Step-by-Step 解説

1ポテンシャルの意味づけ
pot[x]は「$x$の値$-$親の値」をmod $P$で表す。経路圧縮後は「$x$の値$-$根の値」になる不変量を崩さない設計。
2反復findで経路圧縮
再帰findは$N=2\times10^5$規模の鎖状構造で危険。根までの経路を記録し末尾から累積して根に直接張替える。
3unionの代数
$x_{r_b}-x_{r_a}=w+p_a-p_b$を使う。Union by Rankで向きを入れ替えるときは$w$の符号も反転させる。
4矛盾判定
同じ根に属するとき、導出値と申告値がmod $P$で一致すれば無矛盾、しなければ制約全体を無視(構造は変更しない)。
5クエリ処理
根が異なれば「不明」=-1、同じ根なら$(pot[b]-pot[a])\bmod P$を返す。

よくあるミス

ミス原因正しい書き方
符号反転時にmodを取り忘れるw=-wのまま計算し負の値が伝播w=(-w)%Pを反転直後に必ず適用
rank入れ替え後ポテンシャルの入れ替えを忘れるra,rbだけswapしpa,pbを忘れるra,rbとpa,pbを必ずセットでswap
矛盾した制約を無視するはずが構造を書き換えるfindの経路圧縮とunionの併合を混同findの圧縮は常に安全、併合(par[rb]=ra)のみ矛盾時に行わない
非素数の$P$で逆元を使おうとする乗除算が必要と誤解加算・減算・modのみで実装できる(非素数の$P$でも正しく動く)
pot[x]を根からの相対値だと誤解経路圧縮前は親からの相対値であることを見落とすfind実行直後でなければ根からの相対値として使わない

次のステップ

  • 発展: $P$の代わりに実数上の差分制約に変えた通常の重み付きUnion-Find(ABC087 D相当)も同じ枠組みで解けることを確認せよ。
  • 発展2: 制約を後から取り消す(undo)を追加する場合、Union by Rankのみ(経路圧縮なし)にしたロールバック可能DSUの設計を検討せよ。
  • 次回予告: 次回のテーマは今後のセッションで選定

自己評価

自分の回答

気づき・メモ