#include <bits/stdc++.h>
#define pb push_back
#define spc << " " <<
//#define endl "\n"
#define all(x) x.begin(), x.end()
#define ll long long
#define ii pair<int, int>
#define vi vector<int>
#define vii vector<ii>
#define st first
#define nd second
#define mid (l+r)/2
#define inf 1e15
#define MOD 998244353
#define MX 500005
using namespace std;
ll ans;
vii edges[MX], vt[MX];
int tin[MX], tout[MX], par[MX][20], dep[MX];
int tim=0;
ll dp1[MX], dp2[MX], has1[MX], has2[MX];
void dfs(int node, int pa){
tin[node]=++tim;
par[node][0]=pa;
for(int i=1; i<20; i++) par[node][i] = par[par[node][i-1]][i-1];
for(auto p:edges[node]){
if(p.st==pa) continue;
dep[p.st]=dep[node]+p.nd;
dfs(p.st, node);
}
tout[node]=tim;
}
bool anc(int u, int v){
return tin[u]<=tin[v] && tout[u]>=tout[v];
}
int lca(int u, int v){
if(u==v) return u;
if(dep[v]>dep[u]) swap(u,v);
for(int i=19; i>=0; i--) if(!anc(par[u][i], v)) u = par[u][i];
return par[u][0];
}
void dfs2(int node, int pa){
dp1[node]=dp2[node]=inf;
for(auto i:vt[node]){
if(i.st==pa) continue;
dfs2(i.st, node);
dp1[node]=min(dp1[node], dp1[i.st]+i.nd);
dp2[node]=min(dp2[node], dp2[i.st]+i.nd);
}
ans=min(ans, dp1[node]+dp2[node]);
if(has1[node]){
dp1[node]=0;
ans=min(ans, dp2[node]);
}
else if(has2[node]){
ans=min(ans, dp1[node]);
dp2[node]=0;
}
}
void Init(int N, vi A, vi B, vi D){
for(int i=0; i<N-1; i++){
edges[A[i]].pb({B[i], D[i]});
edges[B[i]].pb({A[i], D[i]});
}
dep[0]=0;
dfs(0,0);
for(int i=0; i<N; i++){
has1[i]=has2[i]=0;
}
}
ll Query(int S, vi X, int T, vi Y){
vi ver;
for(auto i:X){
has1[i]=1;
ver.pb(i);
}
for(auto i:Y){
has2[i]=1;
ver.pb(i);
}
sort(all(ver), [&](int a, int b){ return tin[a]<tin[b]; });
for(int i=1; i<S+T; i++){
ver.pb(lca(ver[i-1], ver[i]));
}
sort(all(ver), [&](int a, int b){ return tin[a]<tin[b]; });
vi ver2 = {ver[0]};
for(int i=1; i<ver.size(); i++) if(ver2.back()!=ver[i]) ver2.pb(ver[i]);
swap(ver, ver2);
stack<int> ss; ss.push(ver[0]);
for(int i=1; i<ver.size(); i++){
while(ss.size()>=2 && !anc(ss.top(), ver[i])){
int f=ss.top(); ss.pop();
vt[f].pb({ss.top(), dep[f]-dep[ss.top()]});
vt[ss.top()].pb({f, dep[f]-dep[ss.top()]});
}
ss.push(ver[i]);
}
while(ss.size()>=2){
int f=ss.top(); ss.pop();
vt[f].pb({ss.top(), dep[f]-dep[ss.top()]});
vt[ss.top()].pb({f, dep[f]-dep[ss.top()]});
}
ans=inf;
dfs2(ver[0], ver[0]);
for(auto i:ver){
has1[i]=has2[i]=0;
vt[i].clear();
}
return ans;
}
Compilation message
factories.cpp: In function 'long long int Query(int, std::vector<int>, int, std::vector<int>)':
factories.cpp:91:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
91 | for(int i=1; i<ver.size(); i++) if(ver2.back()!=ver[i]) ver2.pb(ver[i]);
| ~^~~~~~~~~~~
factories.cpp:94:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
94 | for(int i=1; i<ver.size(); i++){
| ~^~~~~~~~~~~
/usr/bin/ld: /tmp/ccIbhmqY.o: in function `main':
grader.cpp:(.text.startup+0x37d): undefined reference to `Init(int, int*, int*, int*)'
/usr/bin/ld: grader.cpp:(.text.startup+0x412): undefined reference to `Query(int, int*, int, int*)'
collect2: error: ld returned 1 exit status
