#include <bits/stdc++.h>
#include "factories.h"
using namespace std;
const int maxn = 500000 + 5;
const int LOG = 20;
const long long INFLL = (long long)1e18;
int n;
vector<pair<int,long long>> adj[maxn];
int h[maxn];
long long dis[maxn];
int st[maxn][LOG];
// centroid decomp
bool delNode[maxn];
int subSize[maxn];
int parCent[maxn]; // parent in centroid-tree for centroid nodes (optional)
vector<int> centroidsOfNode[maxn];
long long bestAtCent[maxn];
// ---------- LCA ----------
void dfs_lca(int root){
// iterative stack to avoid recursion overflow
stack<pair<int,int>> s;
s.push({root, -1});
h[root] = 0;
dis[root] = 0;
st[root][0] = root;
while(!s.empty()){
auto [u, p] = s.top(); s.pop();
for (auto &e : adj[u]){
int v = e.first;
long long w = e.second;
if (v == p) continue;
h[v] = h[u] + 1;
dis[v] = dis[u] + w;
st[v][0] = u;
s.push({v, u});
}
}
// fill binary lifting table
for (int j = 1; j < LOG; ++j){
for (int i = 0; i < n; ++i){
st[i][j] = st[ st[i][j-1] ][j-1];
}
}
}
int lca(int u, int v){
if (h[u] > h[v]) swap(u,v);
int k = h[v] - h[u];
for (int i = 0; i < LOG; ++i) if (k & (1<<i)) v = st[v][i];
if (u == v) return u;
for (int i = LOG-1; i >= 0; --i){
if (st[u][i] != st[v][i]){
u = st[u][i];
v = st[v][i];
}
}
return st[u][0];
}
inline long long length_uv(int u, int v){
int w = lca(u,v);
return dis[u] + dis[v] - 2*dis[w];
}
// ---------- centroid decomposition ----------
void dfs_size(int u, int p){
subSize[u] = 1;
for (auto &e : adj[u]){
int v = e.first;
if (v == p || delNode[v]) continue;
dfs_size(v, u);
subSize[u] += subSize[v];
}
}
int find_centroid(int u, int p, int total){
for (auto &e : adj[u]){
int v = e.first;
if (v == p || delNode[v]) continue;
if (subSize[v] > total/2) return find_centroid(v, u, total);
}
return u;
}
void mark_component_with_centroid(int u, int p, int centroid){
// add centroid into all nodes of this component
centroidsOfNode[u].push_back(centroid);
for (auto &e : adj[u]){
int v = e.first;
if (v == p || delNode[v]) continue;
mark_component_with_centroid(v, u, centroid);
}
}
void decompose(int entry, int parentC){
dfs_size(entry, -1);
int c = find_centroid(entry, -1, subSize[entry]);
parCent[c] = parentC;
// mark all nodes in this component with centroid c
mark_component_with_centroid(c, -1, c);
// now remove centroid and recurse on components
delNode[c] = true;
for (auto &e : adj[c]){
int v = e.first;
if (!delNode[v]){
decompose(v, c);
}
}
// optionally delNode[c] stays true
}
// ---------- interface ----------
void Init(int N, int A[], int B[], int D[]){
n = N;
// clear structures
for (int i = 0; i < n; ++i){
adj[i].clear();
centroidsOfNode[i].clear();
delNode[i] = false;
bestAtCent[i] = INFLL;
for (int j = 0; j < LOG; ++j) st[i][j] = i; // temp init
}
// build edges (2-way)
for (int i = 0; i < n-1; ++i){
int u = A[i], v = B[i];
long long w = D[i];
adj[u].push_back({v,w});
adj[v].push_back({u,w});
}
// LCA preprocess (root at 0). Use iterative DFS to set st, h, dis
dfs_lca(0);
// centroid decomposition
decompose(0, -1);
// bestAtCent already set to INF
}
long long Query(int S, int X[], int T, int Y[]){
vector<int> touched;
touched.reserve((S+T) * 20);
// update centroids with Y nodes
for (int i = 0; i < T; ++i){
int y = Y[i];
for (int c : centroidsOfNode[y]){
long long val = length_uv(y, c);
if (bestAtCent[c] == INFLL) touched.push_back(c);
if (val < bestAtCent[c]) bestAtCent[c] = val;
}
}
long long res = INFLL;
for (int i = 0; i < S; ++i){
int x = X[i];
for (int c : centroidsOfNode[x]){
if (bestAtCent[c] == INFLL) continue;
long long total = bestAtCent[c] + length_uv(x, c);
if (total < res) res = total;
}
}
// reset touched centroids
for (int c : touched) bestAtCent[c] = INFLL;
return res;
}
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |
| # | Verdict | Execution time | Memory | Grader output |
|---|
| Fetching results... |