shogo-library

競技プログラミング用の C++ ライブラリ

View the Project on GitHub shogo314/shogo-library

:heavy_check_mark: test/aoj/dsl-1-b.2.test.cpp

Depends on

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/problems/DSL_1_B"

#include <functional>
#include <iostream>

#include "../../src/data_structures/potentialized-unionfind.hpp"

// ポテンシャルの群をラムダ式 (関数オブジェクト) から作る
int main() {
    int n, q;
    std::cin >> n >> q;
    PotentializedUnionFind uf(n, Group(std::plus<long long>{}, 0LL, [](long long a) { return -a; }));
    while (q--) {
        int t;
        std::cin >> t;
        if (t == 0) {
            int x, y;
            long long z;
            std::cin >> x >> y >> z;
            uf.merge(x, y, z);
        } else {
            int x, y;
            std::cin >> x >> y;
            if (uf.same(x, y))
                std::cout << uf.diff(x, y) << '\n';
            else
                std::cout << '?' << '\n';
        }
    }
}
#line 1 "test/aoj/dsl-1-b.2.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/problems/DSL_1_B"

#include <functional>
#include <iostream>

#line 2 "src/data_structures/potentialized-unionfind.hpp"

/**
 * @file potentialized-unionfind.hpp
 * @brief ポテンシャル付きUnionFind
 */
#include <algorithm>
#include <cassert>
#line 10 "src/data_structures/potentialized-unionfind.hpp"
#include <stack>
#include <utility>
#include <vector>
#line 2 "src/data_structures/unionfind.hpp"

/**
 * @file unionfind.hpp
 * @brief UnionFind
 */

#line 11 "src/data_structures/unionfind.hpp"

/**
 * @brief 無向グラフに対して「辺の追加」、「2頂点が連結かの判定」をする
 */
struct UnionFind {
   protected:
    int _n;
    // 負ならサイズ、非負なら親
    std::vector<int> parent_or_size;

   public:
    /**
     * @brief merge の戻り値
     *
     * if (uf.merge(a, b)) で新しく統合したかを判定でき、
     * auto [r, merged] = uf.merge(a, b); で代表元も受け取れる
     */
    struct MergeResult {
        int leader;   //!< 統合後の代表元
        bool merged;  //!< 新しく統合したか (もともと同じ連結成分なら false)
        explicit operator bool() const { return merged; }
    };

    UnionFind() : _n(0) {}
    explicit UnionFind(int n) : _n(n), parent_or_size(n, -1) {}

    /**
     * @brief 辺(a,b)を足す
     * @return 統合後の代表元と、新しく統合したか
     */
    MergeResult merge(int a, int b) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        int x = leader(a), y = leader(b);
        if (x == y) return {x, false};
        if (-parent_or_size[x] < -parent_or_size[y]) std::swap(x, y);
        parent_or_size[x] += parent_or_size[y];
        parent_or_size[y] = x;
        return {x, true};
    }
    /**
     * @brief 頂点a,bが連結かどうか
     */
    bool same(int a, int b) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        return leader(a) == leader(b);
    }
    /**
     * @brief 頂点aの属する連結成分の代表元
     */
    int leader(int a) {
        assert(0 <= a && a < _n);
        if (parent_or_size[a] < 0) return a;
        int x = a;
        while (parent_or_size[x] >= 0) {
            x = parent_or_size[x];
        }
        while (parent_or_size[a] >= 0) {
            int t = parent_or_size[a];
            parent_or_size[a] = x;
            a = t;
        }
        return x;
    }
    /**
     * @brief 頂点aの属する連結成分のサイズ
     */
    int size(int a) {
        assert(0 <= a && a < _n);
        return -parent_or_size[leader(a)];
    }
    /**
     * @brief グラフを連結成分に分け、その情報を返す
     * @return 「一つの連結成分の頂点番号のリスト」のリスト
     */
    std::vector<std::vector<int>> groups() {
        std::vector<int> leader_buf(_n), group_size(_n);
        for (int i = 0; i < _n; i++) {
            leader_buf[i] = leader(i);
            group_size[leader_buf[i]]++;
        }
        std::vector<std::vector<int>> result(_n);
        for (int i = 0; i < _n; i++) {
            result[i].reserve(group_size[i]);
        }
        for (int i = 0; i < _n; i++) {
            result[leader_buf[i]].push_back(i);
        }
        result.erase(
            std::remove_if(result.begin(), result.end(),
                           [&](const std::vector<int>& v) { return v.empty(); }),
            result.end());
        return result;
    }
};
#line 2 "src/algebra/group.hpp"

/**
 * @file group.hpp
 * @brief 群の定義
 *
 * 群はモノイド (algebra/monoid.hpp) のメンバに加えて次のメンバを持つ型として表す
 * - S inv(const S& a): 逆元
 */

#line 2 "src/algebra/monoid.hpp"

/**
 * @file monoid.hpp
 * @brief モノイドの定義
 *
 * モノイドは次のメンバを持つ型として表す
 * - using S: 要素の型
 * - S op(const S& a, const S& b): 積 (結合的)
 * - S e(): 単位元
 * op, e は static でも非 static でもよい (データ構造はモノイドのオブジェクトを保持して呼び出す)
 */

#line 15 "src/algebra/monoid.hpp"
#include <istream>
#include <limits>
#include <numeric>
#include <ostream>
#include <type_traits>
#line 21 "src/algebra/monoid.hpp"

/**
 * @brief 1次関数 f(x) = a * x + b
 *
 * monoid::Composite の要素や acted_monoid::AffineSum の作用素に使う
 * 例: Affine<mint> f{2, 3}; f(x) は 2x + 3。省略すると恒等写像 {1, 0}
 */
template <typename T>
struct Affine {
    T a = T(1);  //!< 傾き
    T b = T(0);  //!< 切片
    constexpr T operator()(const T& x) const { return a * x + b; }
    friend constexpr bool operator==(const Affine& f, const Affine& g) { return f.a == g.a && f.b == g.b; }
    friend constexpr bool operator!=(const Affine& f, const Affine& g) { return !(f == g); }
    template <class CharT, class Traits>
    friend std::basic_istream<CharT, Traits>& operator>>(std::basic_istream<CharT, Traits>& is, Affine& f) {
        return is >> f.a >> f.b;
    }
    template <class CharT, class Traits>
    friend std::basic_ostream<CharT, Traits>& operator<<(std::basic_ostream<CharT, Traits>& os, const Affine& f) {
        return os << f.a << CharT(' ') << f.b;
    }
};

namespace monoid {

/**
 * @brief 和
 */
template <typename T>
struct Sum {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a + b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 積
 */
template <typename T>
struct Prod {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a * b; }
    static constexpr S e() { return S(1); }
};

/**
 * @brief 最大値 単位元は型の最小値
 */
template <typename T>
struct Max {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::max(a, b); }
    static constexpr S e() { return std::numeric_limits<S>::lowest(); }
};

/**
 * @brief 最小値 単位元は型の最大値
 */
template <typename T>
struct Min {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::min(a, b); }
    static constexpr S e() { return std::numeric_limits<S>::max(); }
};

/**
 * @brief 最小値と最大値の組 {min, max}
 */
template <typename T>
struct MinMax {
    using S = std::pair<T, T>;
    static constexpr S op(const S& a, const S& b) { return {std::min(a.first, b.first), std::max(a.second, b.second)}; }
    static constexpr S e() { return {std::numeric_limits<T>::max(), std::numeric_limits<T>::lowest()}; }
};

/**
 * @brief 排他的論理和
 */
template <typename T>
struct Xor {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a ^ b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 論理積
 */
template <typename T>
struct And {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a & b; }
    static constexpr S e() { return ~S(0); }
};

/**
 * @brief 論理和
 */
template <typename T>
struct Or {
    using S = T;
    static constexpr S op(const S& a, const S& b) { return a | b; }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 最大公約数 単位元は0
 */
template <typename T>
struct Gcd {
    static_assert(std::is_integral_v<T>);
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::gcd(a, b); }
    static constexpr S e() { return S(0); }
};

/**
 * @brief 最小公倍数 単位元は1
 */
template <typename T>
struct Lcm {
    static_assert(std::is_integral_v<T>);
    using S = T;
    static constexpr S op(const S& a, const S& b) { return std::lcm(a, b); }
    static constexpr S e() { return S(1); }
};

/**
 * @brief 1次関数 (Affine) の合成
 * op(f, g) = g∘f (左から順に適用する)
 * 区間 [l, r) の積は f_{r-1}(...f_l(x)) を表す
 */
template <typename T>
struct Composite {
    using S = Affine<T>;
    static constexpr S op(const S& f, const S& g) { return {g.a * f.a, g.a * f.b + g.b}; }
    static constexpr S e() { return {T(1), T(0)}; }
};

/**
 * @brief 括弧列 対応の取れていない ')' と '(' の個数
 */
struct Parenthesis {
    struct S {
        long long close; // 対応の取れていない ')' の個数
        long long open;  // 対応の取れていない '(' の個数
        constexpr S(char c = 0) : close(c == ')'), open(c == '(') {}
        constexpr S(long long close, long long open) : close(close), open(open) {}
        /**
         * @brief 正しい括弧列か
         */
        constexpr bool is_valid() const { return close == 0 && open == 0; }
    };
    static constexpr S op(const S& a, const S& b) {
        if(a.open >= b.close) return S(a.close, a.open - b.close + b.open);
        return S(a.close + b.close - a.open, b.open);
    }
    static constexpr S e() { return S(); }
};

/**
 * @brief 最小値とその位置 同じ値なら左側を優先
 *
 * @tparam T 値の型
 * @tparam U 位置などの付随する値の型
 */
template <typename T, typename U = int>
struct MinWith {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) { return b.first < a.first ? b : a; }
    static constexpr S e() { return {std::numeric_limits<T>::max(), U{}}; }
};

/**
 * @brief 最大値とその位置 同じ値なら左側を優先
 *
 * @tparam T 値の型
 * @tparam U 位置などの付随する値の型
 */
template <typename T, typename U = int>
struct MaxWith {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) { return a.first < b.first ? b : a; }
    static constexpr S e() { return {std::numeric_limits<T>::lowest(), U{}}; }
};

/**
 * @brief 最小値とその個数
 *
 * @tparam T 値の型
 * @tparam U 個数の型
 */
template <typename T, typename U = long long>
struct MinCount {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) {
        if(a.first < b.first) return a;
        if(b.first < a.first) return b;
        return {a.first, a.second + b.second};
    }
    static constexpr S e() { return {std::numeric_limits<T>::max(), U(0)}; }
};

/**
 * @brief 最大値とその個数
 *
 * @tparam T 値の型
 * @tparam U 個数の型
 */
template <typename T, typename U = long long>
struct MaxCount {
    using S = std::pair<T, U>;
    static constexpr S op(const S& a, const S& b) {
        if(a.first < b.first) return b;
        if(b.first < a.first) return a;
        return {a.first, a.second + b.second};
    }
    static constexpr S e() { return {std::numeric_limits<T>::lowest(), U(0)}; }
};

} // namespace monoid

/**
 * @brief 関数オブジェクトと単位元から実行時に作るモノイド
 * ラムダ式を使える
 *
 * 例: SegTree seg(v, Monoid([](int a, int b) { return std::max(a, b); }, -1));
 *
 * @tparam T 要素の型
 * @tparam Op 積の関数オブジェクトの型
 */
template <typename T, typename Op>
struct Monoid {
    using S = T;
    Op op_object;
    S identity;
    Monoid() = default;
    constexpr Monoid(Op op, const S& e) : op_object(std::move(op)), identity(e) {}
    constexpr S op(const S& a, const S& b) const { return op_object(a, b); }
    constexpr S e() const { return identity; }
};
template <typename Op, typename T>
Monoid(Op, T) -> Monoid<T, Op>;
#line 12 "src/algebra/group.hpp"

namespace group {

/**
 * @brief 加法群
 */
template <typename T>
struct Sum : monoid::Sum<T> {
    using S = T;
    static constexpr S inv(const S& a) { return -a; }
};

/**
 * @brief 乗法群 (0を含まない体の元など)
 */
template <typename T>
struct Prod : monoid::Prod<T> {
    using S = T;
    static constexpr S inv(const S& a) { return S(1) / a; }
};

/**
 * @brief 排他的論理和
 */
template <typename T>
struct Xor : monoid::Xor<T> {
    using S = T;
    static constexpr S inv(const S& a) { return a; }
};

} // namespace group

/**
 * @brief 関数オブジェクトと単位元から実行時に作る群
 * ラムダ式を使える
 *
 * @tparam T 要素の型
 * @tparam Op 積の関数オブジェクトの型
 * @tparam Inv 逆元の関数オブジェクトの型
 */
template <typename T, typename Op, typename Inv>
struct Group {
    using S = T;
    Op op_object;
    S identity;
    Inv inv_object;
    Group() = default;
    constexpr Group(Op op, const S& e, Inv inv) : op_object(std::move(op)), identity(e), inv_object(std::move(inv)) {}
    constexpr S op(const S& a, const S& b) const { return op_object(a, b); }
    constexpr S e() const { return identity; }
    constexpr S inv(const S& a) const { return inv_object(a); }
};
template <typename Op, typename T, typename Inv>
Group(Op, T, Inv) -> Group<T, Op, Inv>;
#line 15 "src/data_structures/potentialized-unionfind.hpp"

/**
 * @brief ポテンシャル付きUnionFind
 * @tparam G ポテンシャルの群 (algebra/group.hpp) 非可換でもよい
 */
template <typename G>
class PotentializedUnionFind {
   public:
    using S = typename G::S;
    using MergeResult = UnionFind::MergeResult;

   private:
    G g;
    int _n;
    // 負ならサイズ、非負なら親
    std::vector<int> parent_or_size;
    // 重み
    std::vector<S> diff_weight;

    S op(const S& a, const S& b) const { return g.op(a, b); }
    S e() const { return g.e(); }
    S inv(const S& a) const { return g.inv(a); }

   public:
    /**
     * @param n 要素数
     * @param g ポテンシャルの群
     */
    explicit PotentializedUnionFind(int n = 0, G g = G()) : g(std::move(g)), _n(n), parent_or_size(n, -1), diff_weight(n, e()) {}

    /**
     * @brief 頂点aの属する連結成分の代表元
     */
    int leader(int a) {
        assert(0 <= a && a < _n);
        if (parent_or_size[a] < 0) return a;
        std::stack<int> stk;
        stk.push(a);
        while (parent_or_size[stk.top()] >= 0) {
            stk.push(parent_or_size[stk.top()]);
        }
        const int root = stk.top();
        stk.pop();
        stk.pop();
        while (!stk.empty()) {
            diff_weight[stk.top()] = op(diff_weight[parent_or_size[stk.top()]], diff_weight[stk.top()]);
            parent_or_size[stk.top()] = root;
            stk.pop();
        }
        return root;
    }

    /**
     * @brief a のグループと b のグループを統合する
     * @param w (b のポテンシャル) - (a のポテンシャル)
     * @return 統合後の代表元と、新しく統合したか
     * @note すでに同じ連結成分なら何もしない (w が矛盾していても検出しない)。
     *       矛盾を調べるには、merged が false のときに diff(a, b) == w かを確かめる
     */
    MergeResult merge(int a, int b, S w) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        w = op(weight(a), w);
        w = op(w, inv(weight(b)));
        int x = leader(a), y = leader(b);
        if (x == y) return {x, false};
        if (-parent_or_size[x] < -parent_or_size[y]) {
            std::swap(x, y);
            w = inv(w);
        }
        parent_or_size[x] += parent_or_size[y];
        parent_or_size[y] = x;
        diff_weight[y] = w;
        return {x, true};
    }

    /**
     * @brief 頂点a,bが連結かどうか
     */
    bool same(int a, int b) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        return leader(a) == leader(b);
    }

    /**
     * @brief (b のポテンシャル) - (a のポテンシャル)
     */
    S diff(int a, int b) {
        assert(same(a, b));
        return op(inv(weight(a)), weight(b));
    }

    /**
     * @brief 頂点aの属する連結成分のサイズ
     */
    int size(int a) {
        assert(0 <= a && a < _n);
        return -parent_or_size[leader(a)];
    }

    /**
     * @brief グラフを連結成分に分け、その情報を返す
     * @return 「一つの連結成分の頂点番号のリスト」のリスト
     */
    std::vector<std::vector<int>> groups() {
        std::vector<int> leader_buf(_n), group_size(_n);
        for (int i = 0; i < _n; i++) {
            leader_buf[i] = leader(i);
            group_size[leader_buf[i]]++;
        }
        std::vector<std::vector<int>> result(_n);
        for (int i = 0; i < _n; i++) {
            result[i].reserve(group_size[i]);
        }
        for (int i = 0; i < _n; i++) {
            result[leader_buf[i]].push_back(i);
        }
        result.erase(std::remove_if(result.begin(), result.end(), [&](const std::vector<int>& v) { return v.empty(); }), result.end());
        return result;
    }

   private:
    S weight(int a) {
        leader(a);
        return diff_weight[a];
    }
};

template <typename G>
PotentializedUnionFind(int, G) -> PotentializedUnionFind<G>;
#line 7 "test/aoj/dsl-1-b.2.test.cpp"

// ポテンシャルの群をラムダ式 (関数オブジェクト) から作る
int main() {
    int n, q;
    std::cin >> n >> q;
    PotentializedUnionFind uf(n, Group(std::plus<long long>{}, 0LL, [](long long a) { return -a; }));
    while (q--) {
        int t;
        std::cin >> t;
        if (t == 0) {
            int x, y;
            long long z;
            std::cin >> x >> y >> z;
            uf.merge(x, y, z);
        } else {
            int x, y;
            std::cin >> x >> y;
            if (uf.same(x, y))
                std::cout << uf.diff(x, y) << '\n';
            else
                std::cout << '?' << '\n';
        }
    }
}
Back to top page