test/library_checker/tree/tree_diameter.test.cpp

• View this file on GitHub
• Last update: 2024-10-21 21:40:42+09:00
• Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

• Dijkstra's Algorithm (ダイクストラ法) (lib/graph/dijkstra.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#include <bits/stdc++.h>
using namespace std;

#include "../../../lib/graph/dijkstra.hpp"

vector<vector<array<long long, 2>>> G;

int n, m;

int main(){
    cin >> n;
    G.resize(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});
    }
    vector<long long> cost = dijkstra(G, 0);
    int s, t;
    long long mx1 = -1, mx2 = -1;
    for(int i = 0; i < n; i++){
        if(cost[i] >= mx1){
            mx1 = cost[i];
            s = i;
        }
    }
    vector<long long> cost2 = dijkstra(G, s);
    for(int i = 0; i < n; i++){
        if(cost2[i] >= mx2){
            mx2 = cost2[i];
            t = i;
        }
    }
    pair<long long, vector<pair<int, int>>> p = shortest_path(G, s, t);
    vector<int> e;
    for(auto [u, v] : p.second){
        if(!e.size() || e.back() != u){
            e.push_back(u);
        }
        e.push_back(v);
    }
    cout << p.first << " " << (int) e.size() << "\n";
    for(auto x : e){
        cout << x << " ";
    }
    cout << "\n";
}

#line 1 "test/library_checker/tree/tree_diameter.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"
#include <bits/stdc++.h>
using namespace std;

#line 2 "lib/graph/dijkstra.hpp"

/**
 * @brief Dijkstra's Algorithm (ダイクストラ法)
 * @docs docs/graph/dijkstra.md
 */

template <typename T>
vector<long long> dijkstra(const vector<vector<array<long long, 2>>> &G, T x){
    const long long INF = 0x1fffffffffffffff;
    vector<long long> cost((int) G.size(), INF);
    using P = pair<long long, long long>;
    priority_queue<P, vector<P>, greater<>> q;
    cost[x] = 0;
    q.emplace(0, x);

    while(q.size()){
        auto [c, at] = q.top();
        q.pop();
        if(c > cost[at]) continue;
        for(auto &[to, t] : G[at]){
            if(cost[to] > c + t){
                cost[to] = c + t;
                q.emplace(cost[to], to);
            }
        }
    }
    return cost;
}

pair<long long, vector<pair<int, int>>> shortest_path(const vector<vector<array<long long, 2>>> &G, int s, int t){
    const long long INF = 0x1fffffffffffffff;
    vector<long long> cost((int) G.size(), INF);
    vector<int> par((int) G.size(), -1);
    using P = pair<long long, long long>;
    priority_queue<P, vector<P>, greater<>> q;
    cost[s] = 0;
    par[s] = -1;
    q.emplace(0, s);

    while(q.size()){
        auto [c, at] = q.top();
        q.pop();
        if(c > cost[at]) continue;
        for(auto &[to, t] : G[at]){
            if(cost[to] > c + t){
                par[to] = at;
                cost[to] = c + t;
                q.emplace(cost[to], to);
            }
        }
    }

    if(cost[t] == INF){
        return {-1, {}};
    }
    vector<pair<int, int>> path;
    for(int now = t; par[now] != -1; now = par[now]){
        path.emplace_back(par[now], now);
    }
    reverse(path.begin(), path.end());

    return {cost[t], path};
}
#line 6 "test/library_checker/tree/tree_diameter.test.cpp"

vector<vector<array<long long, 2>>> G;

int n, m;

int main(){
    cin >> n;
    G.resize(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});
    }
    vector<long long> cost = dijkstra(G, 0);
    int s, t;
    long long mx1 = -1, mx2 = -1;
    for(int i = 0; i < n; i++){
        if(cost[i] >= mx1){
            mx1 = cost[i];
            s = i;
        }
    }
    vector<long long> cost2 = dijkstra(G, s);
    for(int i = 0; i < n; i++){
        if(cost2[i] >= mx2){
            mx2 = cost2[i];
            t = i;
        }
    }
    pair<long long, vector<pair<int, int>>> p = shortest_path(G, s, t);
    vector<int> e;
    for(auto [u, v] : p.second){
        if(!e.size() || e.back() != u){
            e.push_back(u);
        }
        e.push_back(v);
    }
    cout << p.first << " " << (int) e.size() << "\n";
    for(auto x : e){
        cout << x << " ";
    }
    cout << "\n";
}

Back to top page