2022-12-08

AtCoder ABC 087 D - People on a Line (ARC 090 D) (水色, 400 点)

重み付き Union-Find が有効活用できる問題！

問題へのリンク

問題概要

$N$ 個の変数 $x_{1}, x_{2}, \dots, x_{N}$ の値を決定したい。

これらの値を決定するための $M$ 本の制約条件がある。このうち $i$ 個めの情報は 3 つの値 $L_{i}, R_{i}, D_{i}$ によって与えられ、

$x_{R_{i}} - x_{L_{i}} = D_{i}$

であるという制約条件を表す。

これら $M$ 本の制約条件には誤りが含まれる可能性がある。 $M$ 本の制約条件をすべて満たす解が存在するかどうかを判定せよ。

制約

$1 \le N \le 10^{5}$
$0 \le M \le 2 \times 10^{5}$

解法 (1)：重み付き Union-Find

重み付き Union-Find を知っていれば、それを使うだけで解ける問題だとわかる！！

重み付き Union-Find は、Union-Find の各ノードに重みを持たせたデータ構造であり、以下のクエリ処理をサポートできる。なお、この表では Union-Find のノード x の重みを $p_{x}$ と表記している。

クエリ	処理内容
`merge(x, y, w)`	$p_{y} = p_{x} + w$ となるように `x` と `y` を併合する
`issame(x, y)`	ノード `x`, `y` が同じグループにいるかどうかを判定する
`diff(x, y)`	ノード `x`, `y` が同じグループにいるとき、 $p_{y} - p_{x}$ の値を返す

詳細は次の記事にて。

qiita.com

具体的には、 $i = 1, 2, \dots, M$ に対して、次の処理を実行することで解ける。

制約条件がそれまでの条件に矛盾しないかを確かめる
- issame(L[i], R[i]) == False のときは、この 2 値には何の制約もないので問題ない
- issame(L[i], R[i]) == True で、diff(L[i], R[i]) != D[i] の場合は「整合しない」と判定できる
矛盾しない場合は、merge(L[i], R[i], D[i]) を実行する

計算量は $O(N + M\alpha(N))$ 。

コード

#include <bits/stdc++.h>
using namespace std;

// 重み付き Union-Find
template<class Abel> struct UnionFind {
    const Abel UNITY_SUM = 0;   // ゼロに相当する値を入れる
    vector<int> par;
    vector<Abel> diff_weight;

    UnionFind() { }
    UnionFind(int n) : 
    par(n, -1), diff_weight(n, UNITY_SUM) {}
    
    int root(int x) {
        if (par[x] < 0) return x;
        else {
            int r = root(par[x]);
            diff_weight[x] += diff_weight[par[x]];
            return par[x] = r;
        }
    }
    
    Abel calc_weight(int x) {
        root(x);
        return diff_weight[x];
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y, Abel w = 0) {
        w += calc_weight(x); w -= calc_weight(y);
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y), w = -w; // merge technique
        par[x] += par[y];
        par[y] = x;
        diff_weight[y] = w;
        return true;
    }
    
    Abel diff(int x, int y) {
        return calc_weight(y) - calc_weight(x);
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    int N, M;
    cin >> N >> M;

    // 重み付き Union-Find を宣言
    UnionFind<long long> uf(N);

    // 制約条件を順に処理する
    bool res = true;
    for (int i = 0; i < M; ++i) {
        int L, R;
        long long D;
        cin >> L >> R >> D;
        --L, --R;

        if (uf.issame(L, R) && uf.diff(L, R) != D) {
            res = false;
        } else {
            uf.merge(L, R, D);
        }
    }
    cout << (res ? "Yes" : "No") << endl;
}

解法 (2)：BFS や DFS で値を埋めて整合性を確認する

二部グラフ判定問題を BFS / DFS で実行する方法に馴染みのある人なら、それを応用することで解けることがピンと来ると思う！

まず、 $M$ 本の制約条件をそれぞれ「辺」とみなしたグラフを考えることにする。このグラフ上で BFS / DFS を実行する。

頂点 $s$ を始点とした BFS / DFS を開始することを考える。このとき $x_{s} = 0$ としておく。 $x$ は相対的な関係なので、始点 $x$ に対する $x_{x}$ の値は好きに決めて良い。

そして、探索中に辺 $e = (u, v)$ を辿ったときに、次の作業を実行する。

$x_{v}$ の値がまだ決まっていないならば、辺 $e$ に関する制約条件から、 $x_{v}$ の値を決定する
$x_{v}$ の値がすでに決まっているならば、その値が辺 $e$ に関する制約条件と矛盾しないかどうかを確認する

BFS でも DFS でも同様にできる。計算量は $O(N + M)$ となる。

コード

実装上は、グラフに辺 $(L_{i}, R_{i})$ (重み $D_{i}$ ) を追加するときに、辺 $(R_{i}, L_{i})$ (重み $-D_{i}$ ) を一緒に追加すると楽になる。

#include <bits/stdc++.h>
using namespace std;

using Edge = pair<int, long long>;  // 重み付き辺 (終点, 重み)
using Graph = vector<vector<Edge>>;  // グラフ

int main() {
    // 入力
    int N, M;
    cin >> N >> M;
    Graph G(N);
    for (int i = 0; i < M; ++i) {
        int L, R;
        long long D;
        cin >> L >> R >> D;
        --L, --R;

        G[L].emplace_back(R, D);
        G[R].emplace_back(L, -D);
    }

    // BFS
    bool res = true;
    vector<bool> seen(N, false);
    vector<long long> x(N, -1);
    queue<int> que;
    for (int s = 0; s < N; ++s) {
        if (seen[s]) continue;

        // 頂点 s から探索開始
        seen[s] = true;
        x[s] = 0;
        que.push(s);
        while (!que.empty()) {
            int v = que.front();
            que.pop();

            for (Edge e : G[v]) {
                if (seen[e.first]) {
                    // 探索済みの場合は整合性を確認する
                    if (x[e.first] - x[v] != e.second) {
                        res = false;
                    }
                } else {
                    // 探索済みでない場合は x の値を定める
                    seen[e.first] = true;
                    x[e.first] = x[v] + e.second;
                    que.push(e.first);
                }
            }
        }
    }
    cout << (res ? "Yes" : "No") << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 087 D - People on a Line (ARC 090 D) (水色, 400 点)

問題概要

制約

解法 (1)：重み付き Union-Find

コード

解法 (2)：BFS や DFS で値を埋めて整合性を確認する

コード