#include "highway.h"
#include <bits/stdc++.h>
using namespace std;
#define int long long
vector<pair<bool,pair<int,int>>> edges;
vector<vector<int>> to;
vector<int32_t> vec;
int len,n,a,b,s,e;
vector<int> subtrees(int u){
vector<int> subtree(n,1);
stack<pair<int,int>> dfs;
dfs.push({u,-1});
int par,now,v;
while(!dfs.empty()){
now=dfs.top().first;
par=dfs.top().second;
dfs.pop();
for(int i=0;i<to[now].size();i++){
if(edges[to[now][i]].first==false){
continue;
}else if(edges[to[now][i]].second.first==now){
v=edges[to[now][i]].second.second;
}else{
v=edges[to[now][i]].second.first;
}
if(v!=par){
dfs.push({v,now});
}
}
if(par!=-1){
subtree[par]+=subtree[now];
}
}
return subtree;
}
int centroid(int u){
int par,now,v;
vector<int> subtree=subtrees(u);
now=u;
while(true){
par=now;
for(int i=0;i<to[now].size();i++){
if(edges[to[now][i]].first==false){
continue;
}else if(edges[to[now][i]].second.first==now){
v=edges[to[now][i]].second.second;
}else{
v=edges[to[now][i]].second.first;
}
if(subtree[v]*2>=subtree[u]){
now=v;
}
}
if(now==par){
break;
}
}
return now;
}
int solve(int u,int dep){
// cout << u << ' ' << dep << '\n';
// for(int i=0;i<edges.size();i++){
// cout << edges[i].first << " \n"[i+1==edges.size()];
// }
int now,par,d,v,l,r;
vector<pair<int,int>> pos;
stack<pair<pair<int,int>,int>> dfs;
dfs.push({{u,-1},0});
while(!dfs.empty()){
now=dfs.top().first.first;
par=dfs.top().first.second;
d=dfs.top().second;
dfs.pop();
// cout << now << ' ' << par << ' ' << d << '\n';
for(int i=0;i<to[now].size();i++){
// cout << i << ' ' << edges[to[now][i]].first << '\n';
if(edges[to[now][i]].first==false){
continue;
}else if(edges[to[now][i]].second.first==now){
v=edges[to[now][i]].second.second;
}else{
v=edges[to[now][i]].second.first;
}
if(d==dep-1){
pos.push_back({v,to[now][i]});
}else{
dfs.push({{v,now},d+1});
}
}
}
// cout << pos.size() << '\n';
l=0;r=pos.size()-1;
while(l!=r){
int m=(l+r)/2;
fill(vec.begin(),vec.end(),0);
for(int i=0;i<=m;i++){
vec[pos[i].second]=1;
}
if(ask(vec)!=len*a){
r=m;
}else{
l=m+1;
}
}
// cout << pos[l].first << '\n';
return pos[l].first;
}
void divide(int u){
int cen=centroid(u),tmp=0,v,cnt,now,par;
// cout << u << ' ' << cen << '\n';
vector<int> subtree=subtrees(cen),use,unuse;
vector<pair<int,int>> order;
for(int i=0;i<to[cen].size();i++){
if(edges[to[cen][i]].first==false){
continue;
}else if(edges[to[cen][i]].second.first==cen){
v=edges[to[cen][i]].second.second;
}else{
v=edges[to[cen][i]].second.first;
}
order.push_back({subtree[v],to[cen][i]});
}
sort(order.begin(),order.end());
fill(vec.begin(),vec.end(),0);
for(int i=order.size()-1;i>=0;i--){
if((tmp+order[i].first)*2<subtree[cen]){
tmp+=order[i].first;
use.push_back(order[i].second);
}else{
unuse.push_back(order[i].second);
edges[unuse.back()].first=false;
}
}
stack<pair<int,int>> dfs;
dfs.push({cen,-1});
while(!dfs.empty()){
now=dfs.top().first;
par=dfs.top().second;
dfs.pop();
for(int i=0;i<to[now].size();i++){
if(edges[to[now][i]].first==false){
continue;
}else if(edges[to[now][i]].second.first==now){
v=edges[to[now][i]].second.second;
}else{
v=edges[to[now][i]].second.first;
}
vec[to[now][i]]=1;
if(v!=par){
dfs.push({v,now});
}
}
}
tmp=ask(vec);
cnt=(tmp-len*a)/(b-a);
if(cnt==len){
divide(cen);
}else if(cnt==0){
for(int i=0;i<use.size();i++){
edges[use[i]].first=false;
}
for(int i=0;i<unuse.size();i++){
edges[unuse[i]].first=true;
}
divide(cen);
}else{
s=solve(cen,cnt);
for(int i=0;i<use.size();i++){
edges[use[i]].first=false;
}
for(int i=0;i<unuse.size();i++){
edges[unuse[i]].first=true;
}
e=solve(cen,len-cnt);
}
}
void find_pair(int32_t N, std::vector<int32_t> U, std::vector<int32_t> V, int32_t A, int32_t B) {
edges.resize(U.size());
to.resize(N);
for(int i=0;i<U.size();i++){
to[U[i]].push_back(i);
to[V[i]].push_back(i);
edges[i]={true,{U[i],V[i]}};
}
vec.resize(U.size(),0);
len=ask(vec)/A;
n=N;a=A;b=B;
divide(0);
answer(s,e);
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |