Test/AtCoder/abc302_h.test.cpp

• View this file on GitHub
• Last update: 2025-02-23 17:42:37+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A
• Link: https://atcoder.jp/contests/abc302/submissions/63090738
• Link: https://atcoder.jp/contests/abc302/tasks/abc302_ex

Depends on

• Src/DataStructure/DisjointSetUnion/UndoableDisjointSetUnion.hpp
• Src/DataStructure/Undoable/UndoableValue.hpp
• Src/DataStructure/Undoable/UndoableVector.hpp
• ioまわりの設定 (Src/Template/IOSetting.hpp)
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A"
// #define PROBLEM "https://atcoder.jp/contests/abc302/tasks/abc302_ex"

#include "../../Src/Template/IOSetting.hpp"
#include "../../Src/DataStructure/DisjointSetUnion/UndoableDisjointSetUnion.hpp"
#include "../../Src/DataStructure/Undoable/UndoableValue.hpp"
#include "../../Src/DataStructure/Undoable/UndoableVector.hpp"

#include <cassert>
#include <iostream>
using namespace zawa;

/*
 * ABC302-Ex Ball Collector
 * https://atcoder.jp/contests/abc302/submissions/63090738
 */

int main() {
#ifdef ATCODER
    SetFastIO();

    int N;
    std::cin >> N;
    std::vector<std::vector<int>> g(N);
    std::vector<int> A(N), B(N);
    for (int i{} ; i < N ; i++) {
        std::cin >> A[i] >> B[i];
        A[i]--; B[i]--;
    }
    for (int _{} ; _ < N - 1 ; _++) {
        int u, v;
        std::cin >> u >> v;
        u--; v--;
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }
    UndoableValue<int> cur{};
    UndoableVector<int> edge(N);
    UndoableDisjointSetUnion dsu(N);
    std::vector<int> ans(N);
    auto dfs{[&](auto dfs, int v, int p) -> void {
        int a{(int)dsu.leader(A[v])}, b{(int)dsu.leader(B[v])};
        if (a == b) {
            int old_val{std::min(edge[a], (int)dsu.size(a))};
            edge.assign(a, edge[a] + 1);
            int new_val{std::min(edge[a], (int)dsu.size(a))};
            cur.assign(cur.value() - old_val + new_val);
            dsu.merge(a, b);
        }
        else {
            int old_val{std::min(edge[a], (int)dsu.size(a)) + std::min(edge[b], (int)dsu.size(b))};
            dsu.merge(a, b);
            int led{(int)dsu.leader(a)};
            edge.assign(led, edge[a] + edge[b] + 1);
            int new_val{std::min(edge[led], (int)dsu.size(led))};
            cur.assign(cur.value() - old_val + new_val);
        }
        ans[v] = cur.value();
        for (auto x : g[v]) if (x != p) {
            dfs(dfs, x, v);
        }
        dsu.undo();
        cur.undo();
        edge.undo();
    }};
    dfs(dfs, 0, -1);
    for (int i{1} ; i < N ; i++) {
        std::cout << ans[i] << (i + 1 == N ? '\n' : ' ');
    }
#else
    std::cout << "Hello World\n";
#endif
}

#line 1 "Test/AtCoder/abc302_h.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A"
// #define PROBLEM "https://atcoder.jp/contests/abc302/tasks/abc302_ex"

#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/UndoableDisjointSetUnion.hpp"

#line 2 "Src/DataStructure/Undoable/UndoableValue.hpp"

#include <cassert>
#include <utility>
#include <vector>

namespace zawa {

template <class T>
class UndoableValue {
public:

    UndoableValue() = default;

    UndoableValue(const T& v) : v_{v} {}

    UndoableValue(T&& v) : v_{std::move(v)} {}

    void assign(const T& v) {
        save();
        v_ = v;
    }

    void assign(T&& v) {
        save();
        v_ = std::move(v);
    }

    const T& value() const {
        return v_;
    }

    void undo() {
        assert(history_.size());
        v_ = history_.back();
        history_.pop_back();
    }

private:
    T v_{};
    std::vector<T> history_;

    inline void save() {
        history_.emplace_back(v_);
    }
};

} // namespace zawa
#line 2 "Src/DataStructure/Undoable/UndoableVector.hpp"

#line 4 "Src/DataStructure/Undoable/UndoableVector.hpp"

#line 8 "Src/DataStructure/Undoable/UndoableVector.hpp"

namespace zawa {

template <class T>
class UndoableVector {
public:
    using Iterator = typename std::vector<T>::const_iterator;

    UndoableVector() = default;

    UndoableVector(usize n) : vec_(n) {}

    UndoableVector(usize n, const T& v) : vec_(n, v) {}

    UndoableVector(const std::vector<T>& vec) : vec_{vec} {}

    UndoableVector(std::vector<T>&& vec) : vec_{std::move(vec)} {}
    
    inline usize size() const {
        return vec_.size();
    }

    Iterator begin() const {
        return vec_.begin();
    }

    Iterator end() const {
        return vec_.end();
    }

    void assign(usize i, const T& v) {
        assert(i < size());
        save(i);
        vec_[i] = v;
    }

    T operator[](usize i) const {
        assert(i < size());
        return vec_[i];
    }

    void undo() {
        assert(history_.size());
        auto [i, v]{history_.back()};
        vec_[i] = v;
        history_.pop_back();
    }

private:
    std::vector<T> vec_; 
    std::vector<std::pair<usize, T>> history_;

    void save(usize i) {
        history_.emplace_back(i, vec_[i]);
    }
};

} // namespace zawa
#line 6 "Src/DataStructure/DisjointSetUnion/UndoableDisjointSetUnion.hpp"

#include <algorithm>
#line 9 "Src/DataStructure/DisjointSetUnion/UndoableDisjointSetUnion.hpp"

namespace zawa {

class UndoableDisjointSetUnion {
public:
    UndoableDisjointSetUnion() = default;

    UndoableDisjointSetUnion(usize n) : n_{n}, comps_{n}, data_(n, -1) {}

    u32 leader(u32 v) const {
        return data_[v] < 0 ? v : leader(data_[v]);
    }

    inline usize size() const noexcept {
        return n_;
    }

    inline usize size(u32 v) const {
        assert(v < size());
        return static_cast<usize>(-data_[leader(v)]);
    }

    inline usize components() const noexcept {
        return comps_.value();
    }

    bool same(u32 u, u32 v) const {
        assert(u < size());
        assert(v < size());
        return leader(u) == leader(v);
    }

    bool merge(u32 u, u32 v) {
        u = leader(u);
        v = leader(v);
        if (u == v) {
            data_.assign(u, data_[u]);
            data_.assign(v, data_[v]);
            comps_.assign(comps_.value());
            return false;
        }
        else {
            if (data_[u] > data_[v]) std::swap(u, v);
            data_.assign(u, data_[u] + data_[v]);
            data_.assign(v, u);
            comps_.assign(comps_.value() - 1);
            return true;
        }
    }

    void undo() {
        data_.undo();
        data_.undo();
        comps_.undo();
    }

    std::vector<std::vector<u32>> enumerate() const {
        std::vector<std::vector<u32>> res(n_);
        for (u32 v{} ; v < n_ ; v++) {
            res[leader(v)].push_back(v);
        }
        res.erase(std::remove_if(res.begin(), res.end(),
                    [](const auto& arr) -> bool { return arr.empty(); }), res.end());
        return res;
    }

private:
    usize n_{};
    UndoableValue<usize> comps_{};
    UndoableVector<i32> data_{};
};

} // namespace zawa
#line 8 "Test/AtCoder/abc302_h.test.cpp"

#line 11 "Test/AtCoder/abc302_h.test.cpp"
using namespace zawa;

/*
 * ABC302-Ex Ball Collector
 * https://atcoder.jp/contests/abc302/submissions/63090738
 */

int main() {
#ifdef ATCODER
    SetFastIO();

    int N;
    std::cin >> N;
    std::vector<std::vector<int>> g(N);
    std::vector<int> A(N), B(N);
    for (int i{} ; i < N ; i++) {
        std::cin >> A[i] >> B[i];
        A[i]--; B[i]--;
    }
    for (int _{} ; _ < N - 1 ; _++) {
        int u, v;
        std::cin >> u >> v;
        u--; v--;
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }
    UndoableValue<int> cur{};
    UndoableVector<int> edge(N);
    UndoableDisjointSetUnion dsu(N);
    std::vector<int> ans(N);
    auto dfs{[&](auto dfs, int v, int p) -> void {
        int a{(int)dsu.leader(A[v])}, b{(int)dsu.leader(B[v])};
        if (a == b) {
            int old_val{std::min(edge[a], (int)dsu.size(a))};
            edge.assign(a, edge[a] + 1);
            int new_val{std::min(edge[a], (int)dsu.size(a))};
            cur.assign(cur.value() - old_val + new_val);
            dsu.merge(a, b);
        }
        else {
            int old_val{std::min(edge[a], (int)dsu.size(a)) + std::min(edge[b], (int)dsu.size(b))};
            dsu.merge(a, b);
            int led{(int)dsu.leader(a)};
            edge.assign(led, edge[a] + edge[b] + 1);
            int new_val{std::min(edge[led], (int)dsu.size(led))};
            cur.assign(cur.value() - old_val + new_val);
        }
        ans[v] = cur.value();
        for (auto x : g[v]) if (x != p) {
            dfs(dfs, x, v);
        }
        dsu.undo();
        cur.undo();
        edge.undo();
    }};
    dfs(dfs, 0, -1);
    for (int i{1} ; i < N ; i++) {
        std::cout << ans[i] << (i + 1 == N ? '\n' : ' ');
    }
#else
    std::cout << "Hello World\n";
#endif
}

Back to top page