test/yosupo/tree_diameter.test.cpp

• View this file on GitHub
• Last update: 2024-01-07 23:25:49+09:00
• Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

• Diameter of a Tree (tree/tree_diameter.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"

#include "../../tree/tree_diameter.hpp"

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

template <typename T>
ostream& operator<<(ostream& os, const vector<T>& v) {
    for (int i = 0; i < (int)v.size(); ++i) cout << v[i] << " ";
    return os;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    vector<vector<pair<int, ll>>> G(N);
    for (int i = 0; i < N - 1; ++i) {
        int a, b, c;
        cin >> a >> b >> c;
        G[a].push_back({b, c});
        G[b].push_back({a, c});
    }
    auto ans = tree_diameter(G);
    cout << ans.first << " " << ans.second.size() << endl;
    cout << ans.second << endl;
}

#line 1 "test/yosupo/tree_diameter.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"

#line 2 "tree/tree_diameter.hpp"
#include <algorithm>
#include <utility>
#include <vector>

std::pair<int, std::vector<int>> tree_diameter(
    const std::vector<std::vector<int>>& G) {
    std::vector<int> to(G.size());

    auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<int, int> {
        std::pair<int, int> ret(0, v);
        for (int c : G[v]) {
            if (c == p) continue;
            auto weight = dfs(dfs, c, v);
            ++weight.first;
            if (ret < weight) {
                ret = weight;
                to[v] = c;
            }
        }
        return ret;
    };

    auto p = dfs(dfs, 0, -1);
    auto q = dfs(dfs, p.second, -1);
    std::vector<int> path;
    int v = p.second;
    while (v != q.second) {
        path.push_back(v);
        v = to[v];
    }
    path.push_back(v);
    return {q.first, path};
}

template <typename T>
std::pair<T, std::vector<int>> tree_diameter(
    const std::vector<std::vector<std::pair<int, T>>>& G) {
    std::vector<int> to(G.size());

    auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<T, int> {
        std::pair<T, int> ret(0, v);
        for (auto& [u, w] : G[v]) {
            if (u == p) continue;
            auto weight = dfs(dfs, u, v);
            weight.first += w;
            if (ret < weight) {
                ret = weight;
                to[v] = u;
            }
        }
        return ret;
    };

    auto p = dfs(dfs, 0, -1);
    auto q = dfs(dfs, p.second, -1);
    std::vector<int> path;
    int v = p.second;
    while (v != q.second) {
        path.push_back(v);
        v = to[v];
    }
    path.push_back(v);
    return {q.first, path};
}
#line 4 "test/yosupo/tree_diameter.test.cpp"

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

template <typename T>
ostream& operator<<(ostream& os, const vector<T>& v) {
    for (int i = 0; i < (int)v.size(); ++i) cout << v[i] << " ";
    return os;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    cin >> N;
    vector<vector<pair<int, ll>>> G(N);
    for (int i = 0; i < N - 1; ++i) {
        int a, b, c;
        cin >> a >> b >> c;
        G[a].push_back({b, c});
        G[b].push_back({a, c});
    }
    auto ans = tree_diameter(G);
    cout << ans.first << " " << ans.second.size() << endl;
    cout << ans.second << endl;
}

Back to top page