Test/Manual/abc267_f.test.cpp

• View this file on GitHub
• Last update: 2023-12-05 07:56:55+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A
• Link: https://atcoder.jp/contests/abc267/submissions/48201106

Depends on

• Src/Graph/Tree/LevelAncestor.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"

/*
 * ABC267-F Exactly K Steps
 * https://atcoder.jp/contests/abc267/submissions/48201106
 */

#include "../../Src/Template/IOSetting.hpp"
#include "../../Src/Graph/Tree/LevelAncestor.hpp"
using namespace zawa;

#include <algorithm>
#include <iostream>
#include <iterator>
#include <utility>
#include <queue>
#include <vector>

void solve() {
    SetFastIO();
    int n; std::cin >> n;
    std::vector g(n, std::vector<u32>{});
    for (u32 _{} ; _ < (u32)n - 1 ; _++) {
        u32 a, b; std::cin >> a >> b;
        a--; b--;
        g[a].emplace_back(b);
        g[b].emplace_back(a);
    }

    auto bfs{[&](int s) -> std::vector<int> {
        std::vector<int> dist(n, n);
        dist[s] = 0;
        std::queue<int> que{ { s } };
        while (que.size()) {
            int v{que.front()};
            que.pop();
            for (auto x : g[v]) if (dist[x] > dist[v] + 1) {
                dist[x] = dist[v] + 1;
                que.emplace(x);
            }
        }
        return dist;
    }};

    auto dist1{bfs(0)};
    u32 L{(u32)std::distance(dist1.begin(), std::max_element(dist1.begin(), dist1.end()))};
    auto dist2{bfs(L)};
    u32 R{(u32)std::distance(dist2.begin(), std::max_element(dist2.begin(), dist2.end()))};

    LevelAncestor laL{g, L}, laR{g, R};
    int q; std::cin >> q;
    for (int _{} ; _ < q ; _++) {
        int u, k; std::cin >> u >> k;
        u--;
        auto vL{laL(u, k)};
        if (LevelAncestor::invalid(vL)) {
            auto vR{laR(u, k)};
            if (LevelAncestor::invalid(vR)) {
                std::cout << -1 << '\n';
            }
            else {
                std::cout << vR + 1 << '\n';
            }
        }
        else {
            std::cout << vL + 1 << '\n';
        }
    }
}

int main() {
#ifdef ATCODER
    solve();
#else
    std::cout << "Hello World" << '\n';
#endif
}

#line 1 "Test/Manual/abc267_f.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/lesson/2/ITP1/1/ITP1_1_A"

/*
 * ABC267-F Exactly K Steps
 * https://atcoder.jp/contests/abc267/submissions/48201106
 */

#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/Graph/Tree/LevelAncestor.hpp"

#line 4 "Src/Graph/Tree/LevelAncestor.hpp"

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

namespace zawa {

class LevelAncestor {
private:
    usize n_{};
    u32 root_{};
    std::vector<std::vector<u32>> tree_;
    std::vector<std::vector<u32>> jump_;
    std::vector<u32> depth_;

    static constexpr u32 invalid_{static_cast<u32>(-1)};
    static constexpr u32 ceilLog2(u32 n) noexcept {
        assert(n or !"empty graph is not allowed");
        return 32 - __builtin_clz(n);
    }

public:
    constexpr usize size() const noexcept {
        return n_;
    }

    u32 log() const noexcept {
        return jump_.size();
    }

    const u32& depth(u32 v) const noexcept {
        assert(v < size());
        return depth_[v];
    }

    static constexpr u32 invalid() noexcept {
        return invalid_;
    }

    static constexpr bool invalid(usize v) noexcept {
        return v == invalid_;
    }

    constexpr u32 root() const noexcept {
        return root_;
    }

private:

    void dfs(u32 v, u32 p) {
        depth_[v] = (invalid(p) ? invalid_ : depth_[p]) + 1;
        jump_[0][v] = (invalid(p) ? v : p);
        for (u32 i{} ; i + 1 < log() ; i++) {
            jump_[i + 1][v] = jump_[i][jump_[i][v]];
        }
        for (auto x : tree_[v]) if (x != p) {
            assert(invalid(depth_[x]) or !"given graph is not tree");
            dfs(x, v);
        }
    }

public:
    LevelAncestor() = default;
    LevelAncestor(usize n, u32 root = 0)
        : n_{n}, root_{root}, tree_(n), jump_(ceilLog2(n), std::vector<u32>(n)), depth_(n, invalid_) {}
    LevelAncestor(const std::vector<std::vector<u32>>& tree, u32 root = 0) 
        : n_{tree.size()}, root_{root}, tree_{tree}, jump_(ceilLog2(n_), std::vector<u32>(n_)), depth_(n_, invalid_) {
        build();
    }

    void addEdge(usize u, usize v) {
        assert(u < size());
        assert(v < size());
        tree_[u].emplace_back(v);
        tree_[v].emplace_back(u);
    }

    void build() {
        dfs(root(), invalid_);
    }

    usize lca(usize u, usize v) const {
        if (depth(u) > depth(v)) std::swap(u, v);
        for (u32 i{log()} ; i-- ; ) {
            if ((depth(v) - depth(u)) & (1u << i)) {
                v = jump_[i][v];
            }
        }
        if (u == v) return u;
        for (u32 i{log()} ; i-- ; ) {
            if (jump_[i][u] != jump_[i][v]) {
                u = jump_[i][u];
                v = jump_[i][v];
            }
        }
        return jump_[0][u];
    }

    u32 operator()(u32 v, u32 up) const {
        assert(v < size());
        if (depth(v) < up) return invalid();
        if (depth(v) == up) return root();
        for (u32 i{log()} ; i-- ; ) {
            if (up & (1u << i)) {
                up ^= (1u << i);
                v = jump_[i][v];
            }
        }
        return v;
    }

};

} // namespace zawa
#line 10 "Test/Manual/abc267_f.test.cpp"
using namespace zawa;

#include <algorithm>
#line 14 "Test/Manual/abc267_f.test.cpp"
#include <iterator>
#line 16 "Test/Manual/abc267_f.test.cpp"
#include <queue>
#line 18 "Test/Manual/abc267_f.test.cpp"

void solve() {
    SetFastIO();
    int n; std::cin >> n;
    std::vector g(n, std::vector<u32>{});
    for (u32 _{} ; _ < (u32)n - 1 ; _++) {
        u32 a, b; std::cin >> a >> b;
        a--; b--;
        g[a].emplace_back(b);
        g[b].emplace_back(a);
    }

    auto bfs{[&](int s) -> std::vector<int> {
        std::vector<int> dist(n, n);
        dist[s] = 0;
        std::queue<int> que{ { s } };
        while (que.size()) {
            int v{que.front()};
            que.pop();
            for (auto x : g[v]) if (dist[x] > dist[v] + 1) {
                dist[x] = dist[v] + 1;
                que.emplace(x);
            }
        }
        return dist;
    }};

    auto dist1{bfs(0)};
    u32 L{(u32)std::distance(dist1.begin(), std::max_element(dist1.begin(), dist1.end()))};
    auto dist2{bfs(L)};
    u32 R{(u32)std::distance(dist2.begin(), std::max_element(dist2.begin(), dist2.end()))};

    LevelAncestor laL{g, L}, laR{g, R};
    int q; std::cin >> q;
    for (int _{} ; _ < q ; _++) {
        int u, k; std::cin >> u >> k;
        u--;
        auto vL{laL(u, k)};
        if (LevelAncestor::invalid(vL)) {
            auto vR{laR(u, k)};
            if (LevelAncestor::invalid(vR)) {
                std::cout << -1 << '\n';
            }
            else {
                std::cout << vR + 1 << '\n';
            }
        }
        else {
            std::cout << vL + 1 << '\n';
        }
    }
}

int main() {
#ifdef ATCODER
    solve();
#else
    std::cout << "Hello World" << '\n';
#endif
}

Back to top page