#include "dreaming.h"
#include <compare>
#include <vector>
#include <ranges>
#include <algorithm>
#include <limits>
#include <functional>
#include <iostream>
using namespace std;
class VertexAndWeight {
public:
int vertex;
int weight;
VertexAndWeight(int vertex_, int weight_) : weight(weight_), vertex(vertex_) {}
weak_ordering operator<=>(const VertexAndWeight &other) const {
return this->weight <=> other.weight;
}
VertexAndWeight plusWeight(int w) const {
return {vertex, weight+w};
}
VertexAndWeight &maxEquals(const VertexAndWeight &other) {
if(other > *this) {
*this = other;
}
return *this;
}
};
vector<vector<VertexAndWeight>> adjl;
const int NO_PARENT = -1;
// The parent of a given vertex, or NO_PARENT if there is no parent.
vector<int> parent;
vector<int> distFromRoot;
pair<int, int> processTreeAndGetRadiusDiameter(int v);
int travelTime(int N, int M, int L, int A[], int B[], int T[]) {
adjl.resize(N);
distFromRoot.resize(N);
ranges::fill(distFromRoot, -1);
parent.resize(N);
ranges::fill(parent, NO_PARENT);
for(int i : views::iota(0,M)) {
adjl[A[i]].push_back({B[i], T[i]});
adjl[B[i]].push_back({A[i], T[i]});
}
vector<int> forestRadii;
vector<int> forestDiameters;
//cerr << "N = " << N << endl;
for(int v : views::iota(0,N)) {
//cerr << v << " A\n";
if(distFromRoot[v] == -1) { // has not been included in any tree yet
auto [r, d] = processTreeAndGetRadiusDiameter(v);
//cerr << v << " " << r << " " << d << " B\n";
forestRadii.push_back(r);
forestDiameters.push_back(d);
}
}
if(forestDiameters.size() == 1) {
return forestDiameters[0];
} else if (forestDiameters.size() == 2) {
// Connect the centroids of the two trees in the forest by an edge.
return max({forestDiameters[0], forestDiameters[1], forestRadii[0] + L + forestRadii[1]});
} else {
// Connect the centroid of the tree with the largest radius
// to each of the centroids of the other trees.
partial_sort(forestRadii.begin(), forestRadii.begin()+3, forestRadii.end(), greater<int>());
return max({
*max_element(forestDiameters.begin(), forestDiameters.end()),
forestRadii[0] + L + forestRadii[1],
forestRadii[1] + 2*L + forestRadii[2]
});
}
}
VertexAndWeight getFarthestVertexAndWeight(int v, int par, int dRoot);
pair<int, int> processTreeAndGetRadiusDiameter(int v) {
//cerr << "Starting tree " << v << endl;
int d1 = getFarthestVertexAndWeight(v, NO_PARENT, 0).vertex;
//cerr << "Intermediate tree " << v << endl;
auto [d2, diameter] = getFarthestVertexAndWeight(d1, NO_PARENT, 0);
int x = d2;
int radius = numeric_limits<int>::max();
while(x != NO_PARENT) { // each x is a potential centroid
radius = min(radius, max(distFromRoot[x], diameter - distFromRoot[x]));
x = parent[x];
}
//cerr << "Ending tree " << v << " " << radius << " " << diameter << endl;
return {radius, diameter};
}
// Implements DFS
VertexAndWeight getFarthestVertexAndWeight(int v, int par, int dRoot) {
//cerr << " DFS " << v << endl;
parent[v] = par;
distFromRoot[v] = dRoot;
VertexAndWeight ans(v, 0);
for(auto child : adjl[v]) if (child.vertex != parent[v]) {
VertexAndWeight farthestForThisChild =
getFarthestVertexAndWeight(child.vertex, v, dRoot + child.weight);
ans.maxEquals(farthestForThisChild.plusWeight(child.weight));
}
return ans;
}
