Submission #145267

#	Time	Username	Problem	Language	Result	Execution time	Memory
145267		TadijaSebez	Unique Cities (JOI19_ho_t5)	C++11	100 / 100	1691 ms	85020 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
using namespace std;
#define pb push_back
const int N=200050;
vector<int> E[N];
int col[N],dep[N],par[N];
void DEP(int u, int p)
{
	dep[u]=dep[p]+1;
	par[u]=p;
	for(int v:E[u]) if(v!=p) DEP(v,u);
}
int cen,den,n,m;
int type[N];
bool on_d[N];
const int M=2*N;
const int L=18;
struct SegmentTree
{
	int mx[M],ls[M],rs[M],lzy[M],tsz,root,n;
	SegmentTree(){}
	void init(){ for(int i=1;i<=tsz;i++) ls[i]=rs[i]=mx[i]=lzy[i]=0;root=tsz=n=0;}
	void Build(int &c, int ss, int se, int a[])
	{
		c=++tsz;lzy[c]=0;
		if(ss==se){ mx[c]=a[ss];return;}
		int mid=ss+se>>1;
		Build(ls[c],ss,mid,a);
		Build(rs[c],mid+1,se,a);
		mx[c]=max(mx[ls[c]],mx[rs[c]]);
	}
	void Build(int _n, int a[]){ n=_n;Build(root,1,n,a);}
	void Set(int c, int ss, int se, int qs, int qe, int f)
	{
		if(qs>qe || qs>se || ss>qe) return;
		if(qs<=ss && qe>=se){ lzy[c]+=f;mx[c]+=f;return;}
		int mid=ss+se>>1;
		Set(ls[c],ss,mid,qs,qe,f);
		Set(rs[c],mid+1,se,qs,qe,f);
		mx[c]=max(mx[ls[c]],mx[rs[c]])+lzy[c];
	}
	void Set(int qs, int qe, int f){ Set(root,1,n,qs,qe,f);}
	int Get(int c, int ss, int se, int qs, int qe)
	{
		if(qs<=ss && qe>=se) return mx[c];
		int mid=ss+se>>1;
		if(qe<=mid) return Get(ls[c],ss,mid,qs,qe)+lzy[c];
		if(qs>mid) return Get(rs[c],mid+1,se,qs,qe)+lzy[c];
		return max(Get(ls[c],ss,mid,qs,qe),Get(rs[c],mid+1,se,qs,qe))+lzy[c];
	}
	int Get(int qs, int qe){ return Get(root,1,n,qs,qe);}
} ST;
int lid[N],rid[N],tid,a[N];
void DFS(int u, int p)
{
	dep[u]=dep[p]+1;
	lid[u]=++tid;
	a[tid]=dep[u];
	for(int v:E[u]) if(v!=p) DFS(v,u);
	rid[u]=tid;
}
int cnt[N],anc[N][L],ans[N],sum;
bool work[N];
int GetAnc(int u, int k)
{
	for(int i=0;i<L;i++) if(k>>i&1) u=anc[u][i];
	return u;
}
void Add(int u)
{
	cnt[col[u]]++;
	if(cnt[col[u]]==1) sum++;
}
void Del(int u)
{
	cnt[col[u]]--;
	if(cnt[col[u]]==0) sum--;
}
void Solve(int u, int p, int t)
{
	if(on_d[u] && type[u]==t) work[u]=1;
	if(work[u])
	{
		for(int v:E[u]) if(v!=p) work[v]=1;
	}
	anc[u][0]=p;
	for(int i=1;i<L;i++) anc[u][i]=anc[anc[u][i-1]][i-1];
	vector<int> add;
	ST.Set(lid[u],rid[u],-2);
	int myd=ST.Get(lid[u],lid[u]);
	int d=ST.Get(lid[u],rid[u])-myd;
    int z=GetAnc(u,d);
    if(z!=0 && ST.Get(lid[z],rid[z])==ST.Get(lid[z],lid[z]))
	{
		if(anc[z][0]) add.pb(anc[z][0]);
	}
	z=anc[z][0];
	if(z!=0 && ST.Get(lid[z],rid[z])==ST.Get(lid[z],lid[z]))
	{
		if(anc[z][0]) add.pb(anc[z][0]);
	}
	for(int v:add) Add(v);

	if(work[u])
	{
		//printf("%i %i :%i\n",u,t,sum);
		ans[u]=sum;
	}

	for(int v:E[u]) if(v!=p) Solve(v,u,t);

	ST.Set(lid[u],rid[u],2);
	for(int v:add) Del(v);
}
void Solve(int u, int t)
{
	for(int i=1;i<=n;i++) work[i]=0;
	ST.init();tid=0;
	DFS(u,0);
	ST.Build(tid,a);
	Solve(u,0,t);
}
int main()
{
	int u,v;
	scanf("%i %i",&n,&m);
	for(int i=1;i<n;i++) scanf("%i %i",&u,&v),E[u].pb(v),E[v].pb(u);
	for(int i=1;i<=n;i++) scanf("%i",&col[i]);
	DEP(1,0);for(int i=1;i<=n;i++) if(dep[i]>dep[cen]) cen=i;
	DEP(cen,0);for(int i=1;i<=n;i++) if(dep[i]>dep[den]) den=i;
	for(int i=den;i;i=par[i]) on_d[i]=1;
	int hal=dep[den]/2;
	for(int i=den;dep[i]>hal;i=par[i]) type[i]=1;
	Solve(cen,1);
	Solve(den,0);
	for(int i=1;i<=n;i++) printf("%i\n",ans[i]);
	return 0;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In member function 'void SegmentTree::Build(int&, int, int, int*)':
joi2019_ho_t5.cpp:27:13: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   int mid=ss+se>>1;
           ~~^~~
joi2019_ho_t5.cpp: In member function 'void SegmentTree::Set(int, int, int, int, int, int)':
joi2019_ho_t5.cpp:37:13: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   int mid=ss+se>>1;
           ~~^~~
joi2019_ho_t5.cpp: In member function 'int SegmentTree::Get(int, int, int, int, int)':
joi2019_ho_t5.cpp:46:13: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   int mid=ss+se>>1;
           ~~^~~
joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:126:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%i %i",&n,&m);
  ~~~~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:127:54: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  for(int i=1;i<n;i++) scanf("%i %i",&u,&v),E[u].pb(v),E[v].pb(u);
                       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~
joi2019_ho_t5.cpp:128:29: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  for(int i=1;i<=n;i++) scanf("%i",&col[i]);
                        ~~~~~^~~~~~~~~~~~~~

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32