Test/AOJ/1330.test.cpp

View this file on GitHub
Last update: 2025-02-23 17:42:37+09:00
Problem: https://onlinejudge.u-aizu.ac.jp/problems/1330

Depends on

Code

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

#include "../../Src/Template/IOSetting.hpp"
#include "../../Src/DataStructure/DisjointSetUnion/PotentializedDisjointSetUnion.hpp"
#include "../../Src/Algebra/Group/AdditiveGroup.hpp"

using namespace zawa;

#include <cassert>
#include <iostream>
int N, Q;
bool solve() {
    std::cin >> N >> Q;
    if (N == 0 and Q == 0) return false;
    PotentializedDisjointSetUnion<AdditiveGroup<int>> dsu(N);
    while (Q--) {
        char t;
        int a, b;
        std::cin >> t >> a >> b;
        a--; b--;
        if (t == '!') {
            int w;
            std::cin >> w;
            dsu.merge(a, b, w);
        }
        else if (t == '?') {
            if (dsu.isDefined(a, b)) std::cout << dsu.distance(a, b) << '\n';
            else std::cout << "UNKNOWN\n";
        }
        else assert(false);
    }
    return true;
}
int main() {
    SetFastIO(); 
    while (solve()) ;
}

#line 1 "Test/AOJ/1330.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/problems/1330"

#line 2 "Src/Template/IOSetting.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Template/IOSetting.hpp"

#include <iostream>
#include <iomanip>

namespace zawa {

void SetFastIO() {
    std::cin.tie(nullptr)->sync_with_stdio(false);
}

void SetPrecision(u32 dig) {
    std::cout << std::fixed << std::setprecision(dig);
}

} // namespace zawa
#line 2 "Src/DataStructure/DisjointSetUnion/PotentializedDisjointSetUnion.hpp"

#line 4 "Src/DataStructure/DisjointSetUnion/PotentializedDisjointSetUnion.hpp"

#include <algorithm>
#include <cassert>
#include <numeric>
#include <vector>

namespace zawa {

template <class Group>
class PotentializedDisjointSetUnion {
public:
    using Value = typename Group::Element;
private:
    usize n_{}, comps_{};
    std::vector<u32> parent_{};
    std::vector<u32> size_{};
    std::vector<Value> potential_{};

    u32 leader(u32 v) {
        if (parent_[v] == v) {
            return v;
        }
        else {
            u32 res{leader(parent_[v])};
            potential_[v] = Group::operation(potential_[parent_[v]], potential_[v]);
            return parent_[v] = res;
        }
    }
    Value potential(u32 v) {
        leader(v);
        return potential_[v];
    }

public:

    PotentializedDisjointSetUnion() = default;

    PotentializedDisjointSetUnion(u32 n) 
        : n_{n}, comps_{n}, parent_(n), size_(n, u32{1}), potential_(n, Group::identity()) {
        std::iota(parent_.begin(), parent_.end(), u32{});
    }

    constexpr u32 size() const noexcept {
        return n_;
    }

    u32 size(u32 v) {
        leader(v);
        return size_[v];
    }

    inline u32 components() const noexcept {
        return comps_;
    }

    bool isDefined(u32 u, u32 v) {
        return leader(u) == leader(v);
    }

    Value distance(u32 u, u32 v) {
        assert(u < size());
        assert(v < size());
        return Group::operation(Group::inverse(potential(u)), potential(v));
    }

    bool merge(u32 u, u32 v, Value value) {
        if (isDefined(u, v)) {
            return distance(u, v) == value;
        }
        comps_--;
        value = Group::operation(potential(u), value);
        value = Group::operation(Group::inverse(potential(v)), value);
        u = leader(u);
        v = leader(v);
        if (size_[u] > size_[v]) {
            value = Group::inverse(value);
            std::swap(u, v);
        }
        size_[u] += size_[v];
        parent_[v] = u;
        potential_[v] = value;
        return true;
    }
};

} // namespace zawa
#line 2 "Src/Algebra/Group/AdditiveGroup.hpp"

namespace zawa {

template <class T>
class AdditiveGroup {
public:
    using Element = T;
    static constexpr T identity() noexcept {
        return T{};
    }
    static constexpr T operation(const T& l, const T& r) noexcept {
        return l + r;
    }
    static constexpr T inverse(const T& v) noexcept {
        return -v;
    }
};

} // namespace zawa
#line 6 "Test/AOJ/1330.test.cpp"

using namespace zawa;

#line 11 "Test/AOJ/1330.test.cpp"
int N, Q;
bool solve() {
    std::cin >> N >> Q;
    if (N == 0 and Q == 0) return false;
    PotentializedDisjointSetUnion<AdditiveGroup<int>> dsu(N);
    while (Q--) {
        char t;
        int a, b;
        std::cin >> t >> a >> b;
        a--; b--;
        if (t == '!') {
            int w;
            std::cin >> w;
            dsu.merge(a, b, w);
        }
        else if (t == '?') {
            if (dsu.isDefined(a, b)) std::cout << dsu.distance(a, b) << '\n';
            else std::cout << "UNKNOWN\n";
        }
        else assert(false);
    }
    return true;
}
int main() {
    SetFastIO(); 
    while (solve()) ;
}