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

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

AtCoder ABC 280 F - Pay or Receive (青色, 500 点)

重み付き Union-Find が使える鮮やかな楽しい問題!

問題概要

頂点数  N (頂点番号が  1, 2, \dots, N) のグラフが与えられる。このグラフには  M 組の辺があり、 i 組目の辺は、

  • 頂点  A_i から頂点  B_i へと、長さ  C_i の有向辺
  • 頂点  B_i から頂点  A_i へと、長さ  -C_i の有向辺

となっている。

このグラフに対して  Q 個のクエリが与えられる。各クエリでは頂点  X, Y が与えられますので、

  • 頂点  X から頂点  Y へと至る最長経路長を答えよ
  • ただし、到達不能な場合は "nan" と出力し、
  • いくらでも大きくできる場合は "inf" と出力せよ。

制約

  •  2 \le N \le 10^{5}
  •  0 \le M \le 10^{5}
  •  1 \le Q \le 10^{5}
  •  0 \le C_{i} \le 10^{9}

考えたこと

もし一般の重み付き有向グラフだと、最悪  Q 個の頂点を始点とした最短経路アルゴリズムを実行する必要があって、ベルマンフォード法を用いると計算量  O(QM \log N) になる。これでは間に合わない。

しかし今回の問題では、「辺の向きが逆だと重みが  -1 倍」という特殊構造になっている。きっとこれを活かすと、高速な解法が生み出されるに違いない。

そう考えて色々試すと、結構厳しい制約を満たしていないと、負閉路が簡単に発生することがわかる。突き詰めると、次のようになっていないといけないことがわかる。


各頂点に値  x_{1}, x_{2}, \dots, x_{N} が存在して、各辺の条件に対して

  •  x_{B} - x_{A} = C_{i}
  •  x_{A} - x_{B} = -C_{i}

を満たす


こうなっていないと負閉路が発生してしまうのだ。ちゃんとした証明は公式解説に。

上のことまでわかると、あとは重み付き Union-Find で実現できることがわかる。

qiita.com

コード

負閉路が存在する連結成分は、負閉路があることを表すフラグを立てることにする。

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

// 重み付き Union-Find
template<class Abel> struct UnionFind {
    const Abel UNITY_SUM = 0;      // to be set
    vector<int> par;
    vector<bool> is_nan;
    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, Q;
    cin >> N >> M >> Q;
    
    // 重み付き Union-Find
    UnionFind<long long> uf(N);

    // その頂点の属する連結成分が inf かどうか
    vector<bool> isinf(N, false);

    // UnionFind の構築
    for (int i = 0; i < M; ++i) {
        long long A, B, X;
        cin >> A >> B >> X;
        --A, --B;

        if (uf.issame(A, B) && uf.diff(A, B) != X) {
            // 辻褄が合わないときは inf を true に
            isinf[uf.root(A)] = true;
        } else {
            uf.merge(A, B, X);
        }
    }

    // クエリ処理
    while (Q--) {
        int X, Y;
        cin >> X >> Y;
        --X, --Y;

        if (!uf.issame(X, Y)) {
            // 連結でない場合
            cout << "nan" << endl;
        } else if (isinf[uf.root(X)]) {
            // 負閉路を含む場合
            cout << "inf" << endl;
        } else {
            cout << uf.diff(X, Y) << endl;
        }
    }
}