#include "dreaming.h"
#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>
using namespace std;
const int N = 1e5+1, INF = 2e9;
int n, m, k;
vector<pair<int, int>> adj[N];
template<int km>
struct kMax {
vector<int> vals;
template<int xkm>
friend kMax<xkm> operator+(kMax<xkm> lhs, const int &rhs) {
for (int i = 0; i < lhs.vals.size(); i++)
lhs.vals[i] += rhs;
return lhs;
}
template<int xkm>
kMax<xkm> operator=(kMax<xkm> rhs) {
this->vals = rhs.vals;
return *this;
}
void insert(int x) {
vals.push_back(x);
sort(vals.begin(), vals.end(), greater<int>());
if (vals.size() > km)
vals.resize(km);
}
void insert(kMax x) {
for (int j : x.vals)
vals.push_back(j);
sort(vals.begin(), vals.end(), greater<int>());
if (vals.size() > km)
vals.resize(km);
}
};
kMax<2> dp[N];
bool vis[N];
int diam = 0, wentroid = 0;
void dfs(int node, int par=-1) {
vis[node] = true;
dp[node].vals = {};
for (const auto [ch, w] : adj[node]) {
if (ch != par) {
dfs(ch, node);
if (adj[ch].size() == 1)
dp[node].insert(w);
else
dp[node].insert(dp[ch] + w);
}
}
if (dp[node].vals.size() == 2)
diam = max(diam, dp[node].vals[0] + dp[node].vals[1]);
else if (dp[node].vals.size())
diam = max(diam, dp[node].vals[0]);
}
void dfs2(int node, int par=-1) {
if (dp[node].vals.size())
wentroid = min(wentroid, dp[node].vals[0]);
else
wentroid = 0;
for (auto [ch, w] : adj[node]) {
if (ch != par) {
auto ndv = dp[node];
auto chv = dp[ch];
if (dp[ch].vals.size() && dp[node].vals[0] == dp[ch].vals[0]+w)
dp[ch].insert((dp[node].vals.size() == 2) ? dp[node].vals[1]+w : w);
else
dp[ch].insert(dp[node].vals[0] + w);
dfs2(ch, node);
dp[node] = ndv;
dp[ch] = chv;
}
}
}
void comp(int node) {
diam = 0, wentroid = INF;
dfs(node);
dfs2(node);
}
int travelTime(int nodes, int edges, int len, int au[], int av[], int aw[]) {
n = nodes, m = edges, k = len;
for (int i = 0; i < m; i++) {
adj[au[i]].push_back({av[i], aw[i]});
adj[av[i]].push_back({au[i], aw[i]});
}
int mxdiam = 0;
kMax<3> ww;
for (int i = 0; i < n; i++) {
if (!vis[i]) {
comp(i);
ww.insert(wentroid);
mxdiam = max(mxdiam, diam);
}
}
if (n == m-1)
return mxdiam;
if (ww.vals.size() == 2)
return max(mxdiam, ww.vals[0] + len + ww.vals[1]);
return max({
mxdiam,
ww.vals[0] + len + ww.vals[1],
ww.vals[1] + 2*len + ww.vals[2]
});
}
