Graph/tests/tree_diameter.test.cpp

• View this file on GitHub
• Last update: 2022-11-21 21:56:30+08:00
• Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

• Graph/tree_diameter.h

Code

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

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

#include "../tree_diameter.h"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)
#define SZ(x) ((int)(x).size())

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n; cin >> n;
    vector<vector<pair<int,int>>> g(n);
    REP(i,n-1) {
        int u, v, cost; cin >> u >> v >> cost;
        g[u].emplace_back(v, cost);
        g[v].emplace_back(u, cost);
    }

    auto [length, path] = tree_diameter(g);
    cout << length << ' ' << SZ(path) << endl;
    for (int x : path) cout << x << ' ';
    cout << endl;
    return 0;
}

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

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

#line 1 "Graph/tree_diameter.h"
// Index from 0
// g should contains both (u, v) and (v, u)
// Return {length, path}
//
// Tested:
// - https://judge.yosupo.jp/problem/tree_diameter
//
// Tree diameter (weighted) {{{
using ll = long long;
pair<ll, vector<int>> tree_diameter(const vector<vector<pair<int,int>>>& g) {
    int n = g.size();
    vector<ll> dist(n);
    vector<int> parent(n);

    function<void(int, int, ll)> dfs = [&] (int u, int fu, ll cur_dist) {
        dist[u] = cur_dist;
        parent[u] = fu;
        for (auto [v, cost] : g[u]) if (v != fu) {
            dfs(v, u, cur_dist + cost);
        }
    };
    dfs(0, -1, 0);
    // r = furthest node from root
    int r = max_element(dist.begin(), dist.end()) - dist.begin();
    dfs(r, -1, 0);
    // r->s = longest path
    int s = max_element(dist.begin(), dist.end()) - dist.begin();

    vector<int> path;
    for (int x = s; x >= 0; x = parent[x]) path.push_back(x);

    return {dist[s], path};
}
// }}}
#line 7 "Graph/tests/tree_diameter.test.cpp"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)
#define SZ(x) ((int)(x).size())

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n; cin >> n;
    vector<vector<pair<int,int>>> g(n);
    REP(i,n-1) {
        int u, v, cost; cin >> u >> v >> cost;
        g[u].emplace_back(v, cost);
        g[v].emplace_back(u, cost);
    }

    auto [length, path] = tree_diameter(g);
    cout << length << ' ' << SZ(path) << endl;
    for (int x : path) cout << x << ' ';
    cout << endl;
    return 0;
}

Back to top page