Submission #1014288

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014288		ttamx	JOI tour (JOI24_joitour)	C++17	100 / 100	568 ms	86100 KiB

joitour

#include "joitour.h"
#include <bits/stdc++.h>

using namespace std;

using ll = long long;

const int N=2e5+5;

int n;
vector<int> adj[N];
int a[N];

enum Type{Compress,Rake,AddEdge,AddVertex,Vertex};

struct StaticTopTree{
	int root,node_id;
	int hv[N],sz[N],p[N];
	int lch[4*N],rch[4*N],par[4*N];
	Type type[4*N];
	void dfs(int u){
		sz[u]=1;
		for(auto v:adj[u])if(v!=p[u]){
			p[v]=u;
			dfs(v);
			sz[u]+=sz[v];
			if(sz[v]>sz[hv[u]])hv[u]=v;
		}
	}
	int add(int i,int l,int r,Type t){
		if(!i)i=++node_id;
		lch[i]=l,rch[i]=r,type[i]=t;
		if(l)par[l]=i;
		if(r)par[r]=i;
		return i;
	}
	pair<int,int> merge(const vector<pair<int,int>> &a,Type t){
		if(a.size()==1)return a[0];
		int tot=0;
		vector<pair<int,int>> b,c;
		for(auto [i,s]:a)tot+=s;
		for(auto [i,s]:a){
			(tot>s?b:c).emplace_back(i,s);
			tot-=s*2;
		}
		auto [i,si]=merge(b,t);
		auto [j,sj]=merge(c,t);
		return {add(0,i,j,t),si+sj};
	}
	pair<int,int> compress(int i){
		vector<pair<int,int>> a{add_vertex(i)};
		while(hv[i])a.emplace_back(add_vertex(i=hv[i]));
		return merge(a,Compress);
	}
	pair<int,int> rake(int i){
		vector<pair<int,int>> a;
		for(auto j:adj[i])if(j!=p[i]&&j!=hv[i])a.emplace_back(add_edge(j));
		return a.empty()?make_pair(0,0):merge(a,Rake);
	}
	pair<int,int> add_edge(int i){
		auto [j,s]=compress(i);
		return {add(0,j,0,AddEdge),s};
	}
	pair<int,int> add_vertex(int i){
		auto [j,s]=rake(i);
		return {add(i,j,0,j?AddVertex:Vertex),s+1};
	}
	void build(){
		node_id=n;
		dfs(1);
		root=compress(1).first;
	}
}stt;

struct Point{
	ll c0,c1,c2,c10,c12,c02;
	ll ans;
}point[4*N];

struct Path{
	ll c0,c1,c2,c10,c12;
	ll p0,p1,p2,p10,p12;
	ll s0,s1,s2;
	ll ans;
}path[4*N];

Path compress(Path p,Path c){
	Path res;
	res.c0=p.c0+c.c0,res.c1=p.c1+c.c1,res.c2=p.c2+c.c2;
	res.p0=p.p0+c.p0,res.p1=p.p1+c.p1,res.p2=p.p2+c.p2;
	res.s0=p.s0+c.s0,res.s1=p.s1+c.s1,res.s2=p.s2+c.s2;
	res.c10=p.c10+c.c10+p.s1*c.c0;
	res.c12=p.c12+c.c12+p.s1*c.c2;
	res.p10=p.p10+c.p10+c.s1*p.p0;
	res.p12=p.p12+c.p12+c.s1*p.p2;
	res.ans=p.ans+c.ans+p.p0*c.c12+p.p2*c.c10+c.c0*p.p12+c.c2*p.p10;
	return res;
}

Point rake(Point l,Point r){
	Point res;
	res.c0=l.c0+r.c0;
	res.c1=l.c1+r.c1;
	res.c2=l.c2+r.c2;
	res.c02=l.c02+r.c02+l.c0*r.c2+l.c2*r.c0;
	res.c10=l.c10+r.c10;
	res.c12=l.c12+r.c12;
	res.ans=l.ans+r.ans+l.c0*r.c12+l.c2*r.c10+r.c0*l.c12+r.c2*l.c10;
	return res;
}

Point add_edge(Path x){
	Point res;
	res.c0=x.c0,res.c1=x.c1,res.c2=x.c2;
	res.c10=x.c10,res.c12=x.c12;
	res.c02=0;
	res.ans=x.ans;
	return res;
}

Path add_vertex(Point x,int i){
	Path res;
	res.s0=a[i]==0,res.s1=a[i]==1,res.s2=a[i]==2;
	res.c0=res.p0=x.c0+res.s0;
	res.c1=res.p1=x.c1+res.s1;
	res.c2=res.p2=x.c2+res.s2;
	res.c10=res.p10=x.c10+res.s1*x.c0;
	res.c12=res.p12=x.c12+res.s1*x.c2;
	res.ans=x.ans+res.s1*x.c02+res.s0*x.c12+res.s2*x.c10;
	return res;
}

Path vertex(int i){
	Path res;
	res.c0=res.p0=res.s0=a[i]==0;
	res.c1=res.p1=res.s1=a[i]==1;
	res.c2=res.p2=res.s2=a[i]==2;
	res.c10=res.c12=res.p10=res.p12=res.ans=0;
	return res;
}

void update(int i){
	if(stt.type[i]==Compress)path[i]=compress(path[stt.lch[i]],path[stt.rch[i]]);
	else if(stt.type[i]==Rake)point[i]=rake(point[stt.lch[i]],point[stt.rch[i]]);
	else if(stt.type[i]==AddEdge)point[i]=add_edge(path[stt.lch[i]]);
	else if(stt.type[i]==AddVertex)path[i]=add_vertex(point[stt.lch[i]],i);
	else if(stt.type[i]==Vertex)path[i]=vertex(i);
}

void dfs(int i){
	if(!i)return;
	dfs(stt.lch[i]);
	dfs(stt.rch[i]);
	update(i);
}

void init(int _n,vector<int> _a,vector<int> _u,vector<int> _v,int Q){
	n=_n;
	for(int i=1;i<=n;i++)a[i]=_a[i-1];
	for(int i=0;i<n-1;i++){
		int u=_u[i]+1,v=_v[i]+1;
		adj[u].emplace_back(v);
		adj[v].emplace_back(u);
	}
	stt.build();
	dfs(stt.root);
}

void change(int i,int t){
	i++;
	a[i]=t;
	for(int u=i;u;u=stt.par[u])update(u);
}

ll num_tours(){
	return path[stt.root].ans;
}

• Subtask 1: 6 / 6
• Subtask 2: 8 / 8
• Subtask 3: 6 / 6
• Subtask 4: 16 / 16
• Subtask 5: 16 / 16
• Subtask 6: 34 / 34
• Subtask 7: 14 / 14