Submission #154325

#	Time	Username	Problem	Language	Result	Execution time	Memory
154325		ryansee	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	1024 ms	87188 KiB

joi2019_ho_t5

#include "bits/stdc++.h"
using namespace std;

#define FAST ios_base::sync_with_stdio(false); cin.tie(0);
#define pb push_back
#define eb emplace_back
#define ins insert
#define ph push
#define f first
#define s second
#define cbr cerr << "hi\n"
#define mmst(x, v) memset((x), v, sizeof ((x)))
#define siz(x) ((ll)x.size())
#define all(x) (x).begin(), (x).end()
#define lbd(x, y) lower_bound(all(x), y)
#define ubd(x, y) upper_bound(all(x), y)
#define btinpct(x) __builtin_popcountll((x))
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());    //can be used by calling rng() or shuffle(A, A+n, rng)
inline long long rand(long long x, long long y) { return (rng() % (y+1-x)) + x; } //inclusivesss
string to_string(char c) {string s(1,c);return s;} template<typename T> inline T gcd(T a,T b){if(a>b)swap(a,b);if(a==0)return b;return gcd(b%a,a);}

#define ll /*long long*/ int 
#define ld long double
#define FOR(ii, ss, ee) for(ll ii = (ss); ii <= (ll)(ee); ++ii)
#define DEC(ii, ss, ee) for(ll ii = (ss); ii >= (ll)(ee); --ii)
typedef pair <ll, ll> pi; typedef pair <pi, ll> spi; typedef pair <pi, pi> dpi;

#define LLINF ((long long) 1e18)//1234567890987654321
#define INF 1234567890ll
// #define cerr if(0)cout
#define MAXN (200006)
ll n, m; pi roots;
vector<int> v[MAXN];
int A[MAXN], p[MAXN][21]; int st[MAXN], en[MAXN], co;
// pi fdist[MAXN], rdist[MAXN];
ll depth[MAXN];
void init(ll x,ll par) { p[x][0]=par; st[x]=co++;
	// rdist[x] = fdist[x] = pi(depth[x], x); 
	// spi mx=spi(pi(depth[x], x), 1);
	// spi fmx=spi(pi(depth[x], x), 1);
	for(auto i:v[x]) if(i^par) {
		depth[i]=depth[x]+1; init(i, x);
		// rdist[x]=max(rdist[x], rdist[i]);
		// if(rdist[i].f > mx.f.f) mx.f=rdist[i], mx.s=1;
		// else if(rdist[i].f == mx.f.f) ++ mx.s;
		// fmx=max(fmx, fdist[i]);
		// if(fdist[i].f > fmx.f.f) fmx.f=fdist[i], fmx.s=1;
		// else if(fdist[i].f == fmx.f.f) ++ fmx.s;
	}
	// ll rmx=-1;
	// for(auto i:v[x]) if(i^par) {
		// if(fdist[i].f != fmx.f.f) rmx=max(rmx, rdist[i].f);
	// }
	// if(mx.s == 1 && fmx.f.f > rmx && fmx.s==1) fdist[x]=fmx.f; 
	en[x]=co;
}
void t2k() { FOR(j,1,20) FOR(i,1,n) p[i][j]=p[p[i][j-1]][j-1]; }
bool isp(ll a,ll b){return st[a]<=st[b]&&en[a]>=en[b];}
ll lca(ll x,ll y) {
	if(isp(x,y))return x; if(isp(y,x))return y;
	for(ll i=20;i>=0;--i) {
		if(!isp(p[x][i],y))x=p[x][i];
	}
	return p[x][0];
}
// ll ans[MAXN]; 
ll get_dist(ll a,ll b) { return depth[a]+depth[b]-2*depth[lca(a,b)]; }
struct diam_tree {
	int root;
	vector<ll> ans;
	vector <pi> dq;
	ll cur=0;
	ll mp[MAXN];
	ll rdist[MAXN], depth[MAXN];
	ll moves=0;
	void add(ll x, ll d) {
	    assert(++moves <= 5e8);
		// cerr<<x<<" in\n";
		if(!mp[x]) ++ cur;
		++ mp[x];
		dq.eb(x, d);
	}
	void del() { assert(dq.size());
		// cerr<<dq.back().f<<" out\n";
		ll x=dq.back().f;
		-- mp[x];
		if(mp[x] == 0) -- cur;
		dq.pop_back();
	}
	void main() {
		ans.resize(n+1, 0); mmst(mp,0); dq.clear(); mmst(depth,0); mmst(rdist,0);
		assert(root);
		cur=0;
		pre(root, root); dfs(root, root); //cerr<<"r: "<<root<<"\n";
		// FOR(i,1,n) cerr<<i<<": "<<ans[i]<<"\n";
	}
	void dfs(ll x,ll p) {
// 		add(A[x], depth[x]);
		// ans[x]=cur;
		pi one, two;
		one=pi(0, -1);
		two=pi(0, -1);
		for(auto i:v[x]) if(i^p) {
		    two=max(two, pi(rdist[i]-depth[x], i));
		    if(two > one) swap(one, two);
		}
		if(one.s != -1) {
			while(dq.size() && dq.back().s >= depth[x]-two.f) del();
			add(A[x], depth[x]);
			dfs(one.s, x);
		}
		for(auto i:v[x]) if((i^p) && (i^one.s)) {
			while(dq.size() && dq.back().s >= depth[x]-one.f) del();
			add(A[x], depth[x]);
			dfs(i, x);
		}
		while(dq.size() && dq.back().s >= depth[x]-one.f) del();
		ans[x]=cur;
	}
	void pre(ll x,ll p) {
		rdist[x]=depth[x];
		for(auto i:v[x]) if(i^p){
			depth[i]=depth[x]+1;
			pre(i,x);
			rdist[x]=max(rdist[x],rdist[i]);
		}
	}
} r1, r2;
bool can(set<pi,greater<pi>> s) {
	if(s.empty()) return 1;
	ll x = s.begin()->f;
	ll times=1;
	s.erase(s.begin());
	for(auto i:s) {
		if(i.f == x) ++ times;
	}
	return times==1;
}
bool decide(ll x){
	if(get_dist(x, roots.f) >= get_dist(x, roots.s)) return 1;
	else return 0;
}
/*void dfs(ll x,ll par,pi dist,pi rup) {
	spi mx = spi(dist, 1);
	ll rmx=0;
	for(auto i:v[x]) if(i^par) {
		fdist[i].f -= depth[x];
		if(fdist[i].f > mx.f.f) mx.f = fdist[i], mx.s=1;
		else if(fdist[i].f == mx.f.f) ++ mx.s;
		fdist[i].f += depth[x];
		// rmx=max(rmx, rdist[i].f-depth[x]);
	}
	if(mx.f.f != dist.f) rmx=rup.f;
	for(auto i:v[x]) if(i^par) if(mx.f.f != fdist[i].f-depth[x]) rmx=max(rmx, rdist[i].f-depth[x]);
	// if(x==3) cerr<<mx.f<<' '<<mx.s<<' '<<rmx<<' '<<dist<<' '<<rup<<"\n";
	if(mx.s == 1 && mx.f.f > rmx) {
		// cerr<<x<<' '<<mx.f.s<<' '<<get_dist(x, mx.f.s)<<' '<<rmx<<' '<<dist.s<<'\n';
		// cerr<<x<<' '<<decide(x)<<'\n';
		if(decide(x)) ans[x]=r1.ans[x];
		else ans[x]=r2.ans[x];
		// ans[x]=1;
	}
	set<spi,greater<spi>> ms;
	multiset<pi,greater<pi>> R;
	ms.ins(spi(dist, x)); R.ins(rup);
	for(auto i:v[x]) if(i^par) {
		ms.ins(spi(pi(fdist[i].f-depth[x], fdist[i].s), i));
		R.ins(pi(rdist[i].f-depth[x],rdist[i].s));
	}
	// while(ms.size() > 4) ms.erase(--ms.end());
	// while(R.size() > 4) R.erase(--R.end());
	ll o=ms.size(); ll oR=R.size();
	for(auto i:v[x])if(i^par){
		ms.erase(spi(pi(fdist[i].f-depth[x], fdist[i].s), i));
		if(ms.begin()->s != x) R.erase(R.find(pi(rdist[ms.begin()->s].f-depth[x],rdist[ms.begin()->s].s))); else R.erase(R.find(rup));
		if(*R.begin() == pi(rdist[i].f-depth[x], rdist[i].s)) R.erase(R.begin());
		if((ms.size() >= 2 && ms.begin()->f.f == next(ms.begin())->f.f) || (R.size()&&ms.begin()->f.f <= R.begin()->f)) {
			if(ms.begin()->s != x) R.ins(pi(rdist[ms.begin()->s].f-depth[x],rdist[ms.begin()->s].s));
			else R.ins(rup);
			dfs(i, x, pi(1,x), pi(R.begin()->f + 1, R.begin()->s));
		} else {
			if(ms.begin()->s != x) R.ins(pi(rdist[ms.begin()->s].f-depth[x],rdist[ms.begin()->s].s));
			else R.ins(rup);
			dfs(i, x, pi((ms.begin()->f.f) + 1, ms.begin()->f.s), pi(R.begin()->f + 1, R.begin()->s));
		}
		if(siz(ms) ^ o) ms.ins(spi(pi(fdist[i].f-depth[x],fdist[i].s),i));
		if(siz(R) ^ oR) R.ins(pi(rdist[i].f-depth[x],rdist[i].s));
	}
}*/
struct diam_find {
	vector<ll> dist;
	pi main() {
		dist.resize(n+1, -1);
		dist[1]=0; dfs(1, 1);
		ll r1=1, r2=-1;
		FOR(i,1,n) if(dist[i] > dist[r1]) r1=i;
		r2=r1;
		dist[r1]=0; dfs(r1, r1);
		FOR(i,1,n) if(dist[i] > dist[r2]) r2=i;
		return pi(r1, r2);
	}
	void dfs(ll x,ll p) {
		for(auto i:v[x]) if(i^p) dist[i]=dist[x]+1, dfs(i, x);
	}
} fd;
int main()
{
	// freopen("int", "r", stdin);
	FAST
	cin>>n>>m;
	FOR(i,0,n-2) {
		ll a, b; cin>>a>>b;
		v[a].eb(b), v[b].eb(a);
	}
	FOR(i,1,n) cin>>A[i];
	roots=fd.main(); init(1, 1); t2k();
	r1.root=roots.f, r2.root=roots.s; r1.main(); r2.main();
	// comp(1);
	FOR(i,1,n) {
		// cerr<<i<<": "<<fdist[i]<<"\n";
	}
	// dfs(1, 1, pi(0,1), pi(0,1));
	FOR(i,1,n) cout<<(decide(i)?r1.ans[i]:r2.ans[i])<<"\n";
}
/*
10 10
2 6
5 8
10 8
1 4
10 6
4 5
10 7
6 9
3 7
1 2 3 4 5 6 7 8 9 10
*/

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'int lca(int, int)':
joi2019_ho_t5.cpp:60:2: warning: this 'if' clause does not guard... [-Wmisleading-indentation]
  if(isp(x,y))return x; if(isp(y,x))return y;
  ^~
joi2019_ho_t5.cpp:60:24: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'if'
  if(isp(x,y))return x; if(isp(y,x))return y;
                        ^~

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