#include "factories.h" #include <iostream> #include <vector> #include <algorithm> #include <queue> using namespace std; const long long inf = 1'000'000'000'000'000; struct idt { int y; long long dist; bool const operator<(const idt &oth) const { return dist>oth.dist; } }; int n,idx; vector<int> nxt,top,sz,cen,active; vector<idt> G; priority_queue<idt> Q[500005]; vector<vector<pair<int,long long>>> up; void add_edge(int x,int y,int dist) { idx++; G[idx]={y,dist}; nxt[idx]=top[x]; top[x]=idx; } int Size(int x,int tata) { sz[x]=1; for(int i=top[x];i;i=nxt[i]) { int y=G[i].y; if(y==tata || cen[y]) continue; sz[x]+=Size(y,x); } return sz[x]; } int centroid(int x,int size,int tata) { for(int i=top[x];i;i=nxt[i]) { int y=G[i].y; if(y==tata || cen[y]) continue; if(2*sz[y]>size) return centroid(y,size,x); } return x; } void dfs(int x,int c,int tata,long long d) { up[x].emplace_back(c,d); for(int i=top[x];i;i=nxt[i]) { auto [y,dist]=G[i]; if(y==tata || cen[y]) continue; dfs(y,c,x,d+dist); } } void decompose(int x,int tata) { int c=centroid(x,Size(x,tata),tata); for(int i=top[c];i;i=nxt[i]) { auto [y,dist]=G[i]; if(y==c || cen[y]) continue; dfs(y,c,c,dist); } cen[c]=1; for(int i=top[c];i;i=nxt[i]) { int y=G[i].y; if(y==c || cen[y]) continue; decompose(y,c); } } void Init(int N,int A[],int B[],int D[]) { n=N; G=vector<idt>(2*n+1); nxt=vector<int>(2*n+1); top=cen=sz=active=vector<int>(n+1); up=vector<vector<pair<int,long long>>>(n+1); for(int i=1;i<=n-1;i++) { add_edge(A[i-1]+1,B[i-1]+1,D[i-1]); add_edge(B[i-1]+1,A[i-1]+1,D[i-1]); } decompose(1,0); } long long dist_to_active(int c) { long long ans=inf; while(!Q[c].empty()) { auto [x,d]=Q[c].top(); if(active[x]) return d; Q[c].pop(); } return ans; } void upd(int x) { Q[x].emplace(x,0); for(auto [c,dist]:up[x]) Q[c].emplace(x,dist); } long long que(int x) { long long ans=inf; for(auto [c,dist]:up[x]) { ans=min(ans,dist+dist_to_active(c)); } return ans; } long long Query(int S,int X[],int T,int Y[]) { long long ans=inf; for(int i=1;i<=S;i++) { upd(X[i-1]+1); active[X[i-1]+1]=1; } for(int i=1;i<=T;i++) { ans=min(ans,que(Y[i-1]+1)); } for(int i=1;i<=S;i++) { active[X[i-1]+1]=0; } return ans; }