Submission #263073

#	Time	Username	Problem	Language	Result	Execution time	Memory
263073		dorijanlendvaj	Mergers (JOI19_mergers)	C++14	100 / 100	1189 ms	164264 KiB

mergers

//DUEL
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#define x first
#define y second
#define pii pair<int,int>
#define pb push_back
#define eb emplace_back
#pragma GCC optimize("unroll-loops")
#define shandom_ruffle(a, b) shuffle(a, b, rng)
#define vi vector<int>
#define vl vector<ll>
#define popcnt __builtin_popcount
#define popcntll __builtin_popcountll
#define all(a) begin(a),end(a)

using namespace std;
using namespace __gnu_pbds;

using ll=long long;
using ull=unsigned long long;
using ld=long double;
int MOD=1000000007;
int MOD2=998244353;
vector<int> bases;
const ll LLINF=1ll<<60;
const char en='\n';

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
void yes() {cout<<"YES"<<en; exit(0);}
void no() {cout<<"NO"<<en; exit(0);}
inline int rund() {int x576363482791fuweh=rng();return abs(x576363482791fuweh)%RAND_MAX;}
template<class T>
void prVec(vector<T> w,bool fl=false)
{
	cout<<w.size()<<en;
	for (int i=0;i<int(w.size())-1;++i) cout<<w[i]<<' ';
	if (w.size()) cout<<w[w.size()-1]<<en;
	if (fl) cout<<flush;
}

void M998()
{
	swap(MOD,MOD2);
}

ll raand()
{
	ll a=rund();
	a*=RAND_MAX;
	a+=rund();
    return a;
}

#define rand raand

ll raaand()
{
	return raand()*(MOD-7)+raand();
}

void compress(vi&v)
{
	set<int> s;
	for (auto a: v) s.insert(a);
	vi o(all(s));
	for (auto&a: v) a=lower_bound(all(o),a)-o.begin();
}

void compress(vl&v)
{
	set<ll> s;
	for (auto a: v) s.insert(a);
	vl o(all(s));
	for (auto&a: v) a=lower_bound(all(o),a)-o.begin();
}

string to_upper(string a)
{
	for (int i=0;i<(int)a.size();++i) if (a[i]>='a' && a[i]<='z') a[i]-='a'-'A';
	return a;
}

string to_lower(string a)
{
	for (int i=0;i<(int)a.size();++i) if (a[i]>='A' && a[i]<='Z') a[i]+='a'-'A';
	return a;
}

ll sti(string a,int base=10)
{
	ll k=0;
	for (int i=0;i<(int)a.size();++i)
	{
		k*=base;
		k+=a[i]-'0';
	}
	return k;
}

template<class T>
void eras(vector<T>& a,T b)
{
	a.erase(find(a.begin(),a.end(),b));
}

string its(ll k,int base=10)
{
	if (k==0) return "0";
	string a;
	while (k)
	{
		a.push_back((k%base)+'0');
		k/=base;
	}
	reverse(a.begin(),a.end());
	return a;
}

ll min(ll a,int b)
{
	if (a<b) return a;
	return b;
}

ll min(int a,ll b)
{
	if (a<b) return a;
	return b;
}

ll max(ll a,int b)
{
	if (a>b) return a;
	return b;
}

ll max(int a,ll b)
{
	if (a>b) return a;
	return b;
}

ll gcd(ll a,ll b)
{
	if (b==0) return a;
	return gcd(b,a%b);
}

ll lcm(ll a,ll b)
{
	return a/gcd(a,b)*b;
}

template<class T,class K>
pair<T,K> mp(T a,K b)
{
	return make_pair(a,b);
}

inline int mult(ll a,ll b)
{
	return (a*b)%MOD;
}

inline int pot(int n,int k)
{
	if (k==0) return 1;
	ll a=pot(n,k/2);
	a=mult(a,a);
	if (k%2) return mult(a,n);
	else return a;
}

int divide(int a,int b)
{
	return mult(a,pot(b,MOD-2));
}

inline int sub(int a,int b)
{
	if (a-b>=0) return a-b;
	return a-b+MOD;
}

inline int add(int a,int b)
{
	if (a+b>=MOD) return a+b-MOD;
	return a+b;
}

bool prime(ll a)
{
	if (a==1) return 0;
	for (int i=2;i<=round(sqrt(a));++i)
	{
		if (a%i==0) return 0;
	}
	return 1;
}

const int N=500010;
int k,n,ind[N],deg[N],c,par[20][N],cn[20][N],pa[N],de[N];
vi im[N],ch[N];

bool cmp(int i,int j)
{
	return ind[i]<ind[j];
}

void dfs(int i,int p)
{
	ind[i]=c++;
	par[0][i]=p;
	for (auto a: ch[i]) if (a!=p) de[a]=de[i]+1,dfs(a,i);
}

void sol(int a,int b)
{
	for (int j=18;j>=0;--j) if (de[a]-de[b]>=(1<<j)) ++cn[j][a],a=par[j][a];
	for (int j=18;j>=0;--j) if (de[b]-de[a]>=(1<<j)) ++cn[j][b],b=par[j][b];
	for (int j=18;j>=0;--j) if (par[j][a]!=par[j][b]) ++cn[j][a],a=par[j][a],++cn[j][b],b=par[j][b];
	if (a!=b)
	{
		int j=0;
		++cn[j][a],a=par[j][a],++cn[j][b],b=par[j][b];
	}
}

int find(int i)
{
	if (i==pa[i]) return i;
	return pa[i]=find(pa[i]);
}

void mer(int a,int b)
{
	a=find(a);
	b=find(b);
	if (a==b) return;
	pa[a]=b;
}

int main()
{
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	for (int i=0;i<10;++i) bases.push_back(rand()%(MOD-13893829*2)[phone]);
	cin>>n>>k;
	for (int i=1;i<n;++i)
	{
		int a,b;
		cin>>a>>b;
		ch[a].pb(b);
		ch[b].pb(a);
	}
	dfs(1,1);
	for (int i=1;i<=n;++i)
	{
		pa[i]=i;
		int z;
		cin>>z;
		im[z].pb(i);
	}
	for (int j=1;j<19;++j) for (int i=1;i<=n;++i) par[j][i]=par[j-1][par[j-1][i]];
	for (int i=1;i<=k;++i)
	{
		sort(all(im[i]),cmp);
		for (int j=1;j<im[i].size();++j) sol(im[i][j-1],im[i][j]);
	}
	for (int j=18;j>=1;--j) for (int i=1;i<=n;++i)
	{
		cn[j-1][i]+=cn[j][i];
		cn[j-1][par[j-1][i]]+=cn[j][i];
	}
	/*for (int i=2;i<=n;++i) cout<<cn[0][i]<<' ';
	cout<<en;*/
	for (int i=2;i<=n;++i) if (cn[0][i]) mer(i,par[0][i]);
	for (int i=2;i<=n;++i) if (!cn[0][i]) ++deg[find(i)],++deg[find(par[0][i])];
	int le=0;
	for (int i=1;i<=n;++i) if (i==find(i) && deg[i]<=1) ++le;
	if (le>1) cout<<(le+1)/2<<en;
	else cout<<0<<en;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:269:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  269 |   for (int j=1;j<im[i].size();++j) sol(im[i][j-1],im[i][j]);
      |                ~^~~~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30