#include<bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef pair<int, int> pii;
typedef vector<pii> vp;
typedef vector<vp> vvp;
vvp graph;
struct DSU {
vi par, subSize;
int size;
DSU() {}
void init(int n) {
size = n;
par.clear();
subSize.clear();
par.resize(n);
subSize.resize(n, 1);
for (int i = 0; i < n; i++) par[i] = i;
}
int findSet(int v) {
if (v == par[v]) return v;
par[v] = findSet(par[v]);
return par[v];
}
bool unite(int v, int u) {
int pv = findSet(v);
int pu = findSet(u);
if (pu == pv) return false;
if (subSize[pu] < subSize[pv]) swap(pv, pu);
par[pv] = pu;
subSize[pu] += subSize[pv];
return true;
}
};
// MST - Minimum spanning tree, Kruskal's algorithm
struct MST {
int cost = 0;
DSU dsu;
int findMST(vector<pair<int, pii>>& edges, int n) {
dsu.init(n);
sort(edges.begin(), edges.end());
int counter = 0;
cost = 0;
for (int i = 0;i < edges.size();i++) {
if (counter == n - 1)break;
int a = edges[i].second.first;
int b = edges[i].second.second;
int w = edges[i].first;
if (dsu.unite(a, b)) cost += w;
}
return cost;
}
};
// LCA - Lowest common ancestor
struct LCA {
vvp table;
vi timeIn;
vi timeOut;
vi depth;
int log;
int timer;
int n;
void dfs(int v, int p, int d, int w) {
depth[v] = d;
timeIn[v] = timer;
table[v][0] = { p, w };
timer++;
for (int i = 0; i < graph[v].size(); i++) {
int u = graph[v][i].first;
int w = graph[v][i].second;
if (u != p) dfs(u, v, d + 1, w);
}
timeOut[v] = timer;
timer++;
}
void preCalc(int N, int root) {
n = N;
log = ceil(log2(n));
timeIn.resize(n + 1);
timeOut.resize(n + 1);
depth.resize(n + 1);
table.resize(n + 1, vp(log + 1));
timer = 0;
dfs(root, root, 0, 0);
for (int i = 1; i <= log; i++) {
for (int v = 1; v <= n; v++) {
int p = table[v][i - 1].first;
int w = table[v][i - 1].second;
table[v][i] = { table[p][i - 1].first, table[p][i - 1].second + w };
}
}
}
bool isAncestor(int a, int b) { // Is a ancestor of b
if (timeIn[a] <= timeIn[b] && timeOut[a] >= timeOut[b]) return true;
return false;
}
int findLCA(int a, int b) {
if (isAncestor(a, b)) return a;
if (isAncestor(b, a)) return b;
for (int i = log; i >= 0; i--) if (!isAncestor(table[a][i].first, b)) a = table[a][i].first;
return table[a][0].first;
}
};
vector<int> shortestPath(int root, int n) {
vector<bool> done(n + 1, false);
priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> pq;
vector<int> dist(n, 1e9);
dist[root] = 0;
pq.push({ 0,root });
while (!pq.empty()) {
int v = pq.top().second;
pq.pop();
if (done[v]) continue;
done[v] = true;
for (int i = 0; i < graph[v].size(); i++) {
int u = graph[v][i].first;
int w = graph[v][i].second;
if (dist[u] > dist[v] + w) {
dist[u] = dist[v] + w;
pq.push({ dist[u],u });
}
}
}
return dist;
}
LCA lca;
void init(int N, vector<int> V, vector<int> U, vector<int> W) {
graph.resize(N);
for (int i = 0;i < N - 1;i++) {
graph[V[i]].push_back({ U[i],W[i] });
graph[U[i]].push_back({ V[i],W[i] });
}
lca.preCalc(N, 0);
}
MST mst;
int query(int a, int b, int c, int d, int e) {
vi nodes = { a,b,c,d,e };
vi newVertices;
for (int i = 0;i < nodes.size();i++) {
for (int j = i;j < nodes.size();j++) {
newVertices.push_back(lca.findLCA(nodes[i], nodes[j]));
}
}
sort(newVertices.begin(), newVertices.end());
newVertices.erase( unique( newVertices.begin(), newVertices.end() ), newVertices.end() );
vvp newGraph(newVertices.size());
vector<pair<int, pii>> edges;
for (int i = 0;i < newVertices.size();i++) {
vector<int> dist = shortestPath(newVertices[i], graph.size());
for (int j = i + 1;j < newVertices.size();j++) {
edges.push_back({ dist[newVertices[j]],{i,j} });
}
}
int ans = mst.findMST(edges, newVertices.size());
return ans;
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, Q;
cin >> N;
vi V(N - 1), U(N - 1), W(N - 1);
for (int i = 0;i < N - 1;i++)cin >> V[i] >> U[i] >> W[i];
init(N, V, U, W);
cin >> Q;
for (int i = 0;i < Q;i++) {
int a, b, c, d, e;
cin >> a >> b >> c >> d >> e;
cout << query(a, b, c, d, e) << endl;
}
return 0;
}
Compilation message
roadsideadverts.cpp: In member function 'int MST::findMST(std::vector<std::pair<int, std::pair<int, int> > >&, int)':
roadsideadverts.cpp:56:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, std::pair<int, int> > >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
56 | for (int i = 0;i < edges.size();i++) {
| ~~^~~~~~~~~~~~~~
roadsideadverts.cpp: In member function 'void LCA::dfs(int, int, int, int)':
roadsideadverts.cpp:85:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
85 | for (int i = 0; i < graph[v].size(); i++) {
| ~~^~~~~~~~~~~~~~~~~
roadsideadverts.cpp: In function 'std::vector<int> shortestPath(int, int)':
roadsideadverts.cpp:142:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
142 | for (int i = 0; i < graph[v].size(); i++) {
| ~~^~~~~~~~~~~~~~~~~
roadsideadverts.cpp: In function 'int query(int, int, int, int, int)':
roadsideadverts.cpp:170:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
170 | for (int i = 0;i < nodes.size();i++) {
| ~~^~~~~~~~~~~~~~
roadsideadverts.cpp:171:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
171 | for (int j = i;j < nodes.size();j++) {
| ~~^~~~~~~~~~~~~~
roadsideadverts.cpp:180:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
180 | for (int i = 0;i < newVertices.size();i++) {
| ~~^~~~~~~~~~~~~~~~~~~~
roadsideadverts.cpp:182:30: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
182 | for (int j = i + 1;j < newVertices.size();j++) {
| ~~^~~~~~~~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
348 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1053 ms |
14940 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1040 ms |
14168 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
348 KB |
Output is correct |
| 2 |
Execution timed out |
1053 ms |
14940 KB |
Time limit exceeded |
| 3 |
Halted |
0 ms |
0 KB |
- |
