Submission #978810

#	Time	Username	Problem	Language	Result	Execution time	Memory
978810		MilosMilutinovic	JOI tour (JOI24_joitour)	C++17	100 / 100	2307 ms	425724 KiB

joitour

#include "joitour.h"

#define pb push_back
#define fi first
#define se second
#define mp make_pair
 
using namespace std;
 
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef long double ld;
 
template <typename T> bool chkmin(T &x,T y){return x>y?x=y,1:0;}
template <typename T> bool chkmax(T &x,T y){return x<y?x=y,1:0;}

int n,tot,ct;
int a[200005],v[400005],nxt[400005],h[200005],sz[200005],dfn[200005][20],dfo[200005][20],c[200005][3],cnt[200005][20][3],dep[200005],cf[200005][20],cr[200005][20];
bool vis[200005];
ll ans,cc[200005][3],ca[200005],pa[200005][20][3],sp[200005][3],pd[200005];

void addedge(int x,int y){
	v[++tot]=y; nxt[tot]=h[x]; h[x]=tot;
	v[++tot]=x; nxt[tot]=h[y]; h[y]=tot;
}

struct fenwick{
	int val[8000005];
	fenwick(){
		for(int i=1;i<=8000000;i++) val[i]=0;
	}
	void change(int p,int v){
		for(;p<=8000000;p+=p&-p) val[p]+=v;
	}
	int ask(int p){
		int v=0;
		for(;p;p-=p&-p) v+=val[p];
		return v;
	}
	int ask(int l,int r){
		return ask(r)-ask(l-1);
	}
}sub[3],up[3];

void dfs1(int x,int f){
	sz[x]=1;
	for(int p=h[x];p;p=nxt[p]){
		if(v[p]==f||vis[v[p]]) continue;
		dfs1(v[p],x);
		sz[x]+=sz[v[p]];
	}
}

int dfs2(int x,int f,int k){
	for(int p=h[x];p;p=nxt[p]){
		if(vis[v[p]]||v[p]==f||sz[v[p]]*2<k) continue;
		return dfs2(v[p],x,k);
	}
	return x;
}

void dfs3(int x,int f,int d,int rt,int r){
	cf[x][d]=rt;
	cr[x][d]=r;
	c[rt][a[x]]++;
	cnt[r][d][a[x]]++;
	dfn[x][d]=++ct;
	for(int p=h[x];p;p=nxt[p]){
		if(v[p]==f||vis[v[p]]) continue;
		dfs3(v[p],x,d,rt,r);
	}
	dfo[x][d]=++ct;
	if(a[x]==1){
		pa[r][d][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
		sp[rt][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
		pa[r][d][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
		sp[rt][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
	}
	up[a[x]].change(dfn[x][d],1);
	up[a[x]].change(dfo[x][d],-1);
	sub[a[x]].change(dfn[x][d],1);
}

void build(int x,int d){
	dfs1(x,x);
	x=dfs2(x,x,sz[x]);
	vis[x]=true;
	cf[x][d]=x;
	dep[x]=d;
	dfn[x][d]=++ct;
	for(int p=h[x];p;p=nxt[p]){
		if(vis[v[p]]) continue;
		dfs3(v[p],x,d,x,v[p]);
	}
	for(int p=h[x];p;p=nxt[p]){
		if(vis[v[p]]) continue;
		ans+=pa[v[p]][d][0]*1ll*(c[x][2]-cnt[v[p]][d][2]);
		ans+=pa[v[p]][d][2]*1ll*(c[x][0]-cnt[v[p]][d][0]);
		pd[x]+=(c[x][0]-cnt[v[p]][d][0])*1ll*cnt[v[p]][d][2];
		pd[x]+=(c[x][2]-cnt[v[p]][d][2])*1ll*cnt[v[p]][d][0];
	}
	dfo[x][d]=++ct;
	if(a[x]==0) ans+=sp[x][2];
	if(a[x]==1) ans+=pd[x]/2;
	if(a[x]==2) ans+=sp[x][0];
	for(int p=h[x];p;p=nxt[p]) if(!vis[v[p]]) build(v[p],d+1);
}

void init(int n,vector<int> f,vector<int> u,vector<int> v,int q){
	::n=n;
	for(int i=0;i+1<n;i++) addedge(u[i],v[i]);
	for(int i=0;i<n;i++) a[i]=f[i];
	build(0,0);
}

void change(int x,int y){
	if(a[x]==y) return;
	for(int d=0;d<=dep[x];d++){
		if(cf[x][d]==x){
			if(a[x]==0) ans-=sp[x][2];
			if(a[x]==1) ans-=pd[x]/2;
			if(a[x]==2) ans-=sp[x][0];
			if(y==0) ans+=sp[x][2];
			if(y==1) ans+=pd[x]/2;
			if(y==2) ans+=sp[x][0];
			continue;
		}
		if(a[cf[x][d]]==0) ans-=sp[cf[x][d]][2];
		if(a[cf[x][d]]==1) ans-=pd[cf[x][d]]/2;
		if(a[cf[x][d]]==2) ans-=sp[cf[x][d]][0];
		if(a[x]==1){
			ans-=sub[0].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][2]-cnt[cr[x][d]][d][2]);
			ans-=sub[2].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][0]-cnt[cr[x][d]][d][0]);
			sp[cf[x][d]][0]-=sub[0].ask(dfn[x][d]+1,dfo[x][d]);	
			sp[cf[x][d]][2]-=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
			pa[cr[x][d]][d][0]-=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
			pa[cr[x][d]][d][2]-=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
		}else{
			ans-=up[1].ask(dfn[x][d])*1ll*(c[cf[x][d]][2-a[x]]-cnt[cr[x][d]][d][2-a[x]]);
			ans-=sp[cf[x][d]][2-a[x]]-pa[cr[x][d]][d][2-a[x]];
			sp[cf[x][d]][a[x]]-=up[1].ask(dfn[x][d]);
			pa[cr[x][d]][d][a[x]]-=up[1].ask(dfn[x][d]);
			pd[cf[x][d]]-=(c[cf[x][d]][2-a[x]]-cnt[cr[x][d]][d][2-a[x]])*2;
		}
		c[cf[x][d]][a[x]]-=1;
		cnt[cr[x][d]][d][a[x]]-=1;
		sub[a[x]].change(dfn[x][d],-1);
		up[a[x]].change(dfn[x][d],-1);
		up[a[x]].change(dfo[x][d],+1);
		if(y==1){
			ans+=sub[0].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][2]-cnt[cr[x][d]][d][2]);
			ans+=sub[2].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][0]-cnt[cr[x][d]][d][0]);
			sp[cf[x][d]][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);	
			sp[cf[x][d]][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
			pa[cr[x][d]][d][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
			pa[cr[x][d]][d][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
		}else{
			ans+=up[1].ask(dfn[x][d])*1ll*(c[cf[x][d]][2-y]-cnt[cr[x][d]][d][2-y]);
			ans+=sp[cf[x][d]][2-y]-pa[cr[x][d]][d][2-y];
			sp[cf[x][d]][y]+=up[1].ask(dfn[x][d]);
			pa[cr[x][d]][d][y]+=up[1].ask(dfn[x][d]);
			pd[cf[x][d]]+=(c[cf[x][d]][2-y]-cnt[cr[x][d]][d][2-y])*2;
		}
		c[cf[x][d]][y]+=1;
		cnt[cr[x][d]][d][y]+=1;
		sub[y].change(dfn[x][d],+1);
		up[y].change(dfn[x][d],+1);
		up[y].change(dfo[x][d],-1);
		if(a[cf[x][d]]==0) ans+=sp[cf[x][d]][2];
		if(a[cf[x][d]]==1) ans+=pd[cf[x][d]]/2;
		if(a[cf[x][d]]==2) ans+=sp[cf[x][d]][0];
	}
	a[x]=y;
}

ll num_tours(){return 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