#include "joitour.h"
#include<bits/stdc++.h>
using namespace std;
vector<int>adj[200005];
int type[200005];
int n;
struct sttttree{
int l[800005],r[800005],p[800005],type[800005],hv[200005],pa[200005],sz[200005];
int cur=0,rt=0;
void dfs(int u,int p){
sz[u]=1;
pa[u]=p;
for(auto x:adj[u])if(x!=p){
dfs(x,u);
if(sz[x]>sz[hv[u]])hv[u]=x;
sz[u]+=sz[x];
}
}
void init(int n){
dfs(1,0);
cur=n;
//cerr<<"in tree:\n";
rt=compress(1).first;
}
int add(int id,int lch,int rch,int t){
if(id==0)id=++cur;
if(lch)p[lch]=id;
if(rch)p[rch]=id;
l[id]=lch,r[id]=rch,type[id]=t;
return id;
}
pair<int,int> merge(int t,vector<pair<int,int>>&v){
if(v.size()==1)return v[0];
vector<pair<int,int>>l,r;
int sz=0;
for(auto x:v)sz+=x.second;
for(auto x:v){
if(sz>=x.second)l.push_back(x);
else r.push_back(x);
sz-=2*x.second;
}
auto lch=merge(t,l);
auto rch=merge(t,r);
int id=add(0,lch.first,rch.first,t);
return {id,lch.second+rch.second};
}
pair<int,int> compress(int x){
//cerr<<"compress "<<x<<"\n";
vector<pair<int,int>>v;
for(x;x!=0;x=hv[x])v.push_back(add_vertex(x));
return merge(1,v);
}
pair<int,int> add_vertex(int x){
//cerr<<"add_vertex "<<x<<"\n";
auto v=rake(x);
return {add(x,v.first,0,v.second==0?3:2),v.second+1};
}
pair<int,int> rake(int x){
//cerr<<"rake "<<x<<"\n";
vector<pair<int,int>>v;
for(auto u:adj[x])if(u!=pa[x]&&u!=hv[x])v.push_back({add_edge(u)});
return v.size()==0?make_pair(0,0):merge(4,v);
}
pair<int,int> add_edge(int x){
//cerr<<"add_edge "<<x<<"\n";
auto v=compress(x);
return {add(0,v.first,0,5),v.second};
}
}tr;
struct path{
long long chain[3];
long long any[3];
//0 1, 2 1
long long pre[2];
//1 0, 1 2
long long suf[2];
long long ans;
path(){
chain[0]=chain[1]=chain[2]=0;
any[0]=any[1]=any[2]=0;
pre[0]=pre[1]=0;
suf[0]=suf[1]=0;
ans=0;
}
};
path paths[800005];
path compress(path a,path b){
path c;
for(int i=0;i<3;i++)c.chain[i]=a.chain[i]+b.chain[i],c.any[i]=a.any[i]+b.any[i];
c.suf[0]=a.chain[1]*b.any[0]+a.suf[0]+b.suf[0];
c.suf[1]=a.chain[1]*b.any[2]+a.suf[1]+b.suf[1];
c.pre[0]=b.chain[1]*a.any[0]+a.pre[0]+b.pre[0];
c.pre[1]=b.chain[1]*a.any[2]+a.pre[1]+b.pre[1];
c.ans=a.ans+b.ans+a.pre[0]*b.any[2]+a.pre[1]*b.any[0]+b.suf[1]*a.any[0]+b.suf[0]*a.any[2];
return c;
}
path add_vertex(path a,int x){
path c;
for(int i=0;i<3;i++)c.any[i]=a.any[i];
c.chain[x]=1;
c.any[x]++;
c.ans=a.ans;
if(x==1){
c.pre[0]=a.any[0];
c.pre[1]=a.any[2];
c.suf[0]=a.any[0];
c.suf[1]=a.any[2];
}
return c;
}
path vertex(int a){
path x;
x.any[a]=1,x.chain[a]=1;
return x;
}
path rake(path a,path b){
path c;
for(int i=0;i<3;i++)c.any[i]=a.any[i]+b.any[i];
c.suf[0]=a.suf[0]+b.suf[0];
c.suf[1]=a.suf[1]+b.suf[1];
c.pre[0]=a.pre[0]+b.pre[0];
c.pre[1]=a.pre[1]+b.pre[1];
c.ans=a.ans+b.ans+b.suf[1]*a.any[0]+b.suf[0]*a.any[2];
return c;
}
path add_edge(path a){
return a;
}
void upd(int x){
if(tr.type[x]==1)paths[x]=compress(paths[tr.l[x]],paths[tr.r[x]]);
else if(tr.type[x]==2)paths[x]=add_vertex(paths[tr.l[x]],type[x]);
else if(tr.type[x]==3)paths[x]=vertex(type[x]);
else if(tr.type[x]==4)paths[x]=rake(paths[tr.l[x]],paths[tr.r[x]]);
else if(tr.type[x]==5)paths[x]=add_edge(paths[tr.l[x]]);
//cerr<<"UPD:"<<x<<"\n";
//cerr<<"ans:"<<paths[x].ans<<"\n";
//cerr<<"one:"<<paths[x].[0]<<' '<<paths[x].one[1]<<" "<<paths[x].one[2]<<"\n";
//cerr<<"two:\n";
/*for(int i=0;i<3;i++){
for(int j=0;j<3;j++)cerr<<paths[x].two[i][j]<<" ";
cerr<<"\n";
}*/
//cerr<<"\n";
}
void dfs(int u){
//cerr<<u<<":"<<tr.type[u]<<"\n";
if(tr.l[u])dfs(tr.l[u]);
if(tr.r[u])dfs(tr.r[u]);
upd(u);
}
void init(int N,vector<int> F,vector<int> U,vector<int> V,int Q) {
n=N;
for(int i=0;i<N-1;i++){
adj[U[i]+1].push_back(V[i]+1);
adj[V[i]+1].push_back(U[i]+1);
//cerr<<U[i]+1<<" "<<V[i]+1<<"\n";
}
for(int i=1;i<=n;i++)type[i]=F[i-1];
//cerr<<"work\n";
tr.init(n);
//cerr<<"work\n";
//cerr<<"\n\n";
dfs(tr.rt);
}
void change(int X, int Y) {
type[++X]=Y;
for(;X!=0;X=tr.p[X])upd(X);
}
long long num_tours() {
return paths[tr.rt].ans;
}
# | 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... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |