제출 #1180748

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1180748		jerzyk	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	389 ms	42400 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;
#define pb push_back
#define st first
#define nd second
typedef long long ll;
typedef long double ld;
const int II = 1'000'000'000;
const int N = 1<<18;
int tab[N], ansc[N], answer[N];

vector<int> ed[N];
int oj[N];
int ds = -1, sa, sb;
int wys[N], sk[N];
vector<int> pth;

bool czy[N];


int lst[N];
int drz[2 * N], sum[2 * N], laz[2 * N];

int clo[N];

void BFS(int n, int s1, int s2)
{
	int v;
	queue<int> q;
	for(int i = 1; i <= n; ++i)
		clo[i] = 0;
	clo[s1] = 2; clo[s2] = 1;
	q.push(s1); q.push(s2);
	while(q.size() > 0)
	{
		v = q.front(); q.pop();
		for(int i = 0; i < (int)ed[v].size(); ++i)
			if(clo[ed[v][i]] == 0)
			{
				clo[ed[v][i]] = clo[v];
				q.push(ed[v][i]);
			}
	}
}

void Change(int v, int p, int k, int pz, int kz, int r)
{
	if(p > kz || k < pz) return;
	if(p >= pz && k <= kz)
	{
		laz[v] += r;
		if(laz[v] == 0)
			drz[v] = sum[v];
		else
			drz[v] = 0;
		return;
	}
	Change(v * 2, p, (p + k) / 2, pz, kz, r);
	Change(v * 2 + 1, (p + k) / 2 + 1, k, pz, kz, r);
	sum[v] = drz[v * 2] + drz[v * 2 + 1];
	if(laz[v] > 0)
		drz[v] = 0;
	else
		drz[v] = sum[v];
}

int Query(int v, int p, int k, int pz, int kz)
{
	if(laz[v] > 0 || p > kz || k < pz) return 0;
	if(p >= pz && k <= kz) return drz[v];
	int a = Query(v * 2, p, (p + k) / 2, pz, kz);
	a += Query(v * 2 + 1, (p + k) / 2 + 1, k, pz, kz);
	return a;
}

void Add(int v, int x)
{
	v += N;
	sum[v] += x;
	if(laz[v] == 0) drz[v] += x;
	v /= 2;
	while(v > 0)
	{
		sum[v] = drz[v * 2] + drz[v * 2 + 1];
		if(laz[v] > 0)
			drz[v] = 0;
		else
			drz[v] = sum[v];
		v /= 2;
	}
}

bool Check(int v)
{
	v += N;
	while(v > 0)
	{
		if(laz[v] > 0) return 1;
		v /= 2;
	}
	return 0;
}

void DFS(int v)
{
	//cout << "DFS: " << v << " " << oj[v] << "\n";
	sk[v] = v;
	pair<int, int> m1 = make_pair(0, v), m2;
	m2 = m1;
	for(int i = 0; i < (int)ed[v].size(); ++i)
	{
		if(oj[ed[v][i]]) continue;
		oj[ed[v][i]] = v;
		DFS(ed[v][i]);
		if(wys[ed[v][i]] + 1 > wys[v])
			sk[v] = sk[ed[v][i]];
		wys[v] = max(wys[v], wys[ed[v][i]] + 1);
		pair<int, int> a = make_pair(wys[ed[v][i]] + 1, sk[ed[v][i]]);
		if(m1 < a) swap(m1, a);
		if(m2 < a) swap(m2, a);
	}
	//cout << v << " " << m1.st << " " << m1.nd << " " << m2.st << " " << m2.nd << "\n";
	if(m1.st + m2.st > ds)
	{
		ds = m1.st + m2.st;
		sa = m1.nd; sb = m2.nd;
	}
	oj[v] = 0;
}

void DFSW(int v)
{
	wys[v] = 0; oj[v] = 1;
	for(int i = 0; i < (int)ed[v].size(); ++i)
	{
		if(oj[ed[v][i]]) continue;
		DFSW(ed[v][i]);
		wys[v] = max(wys[v], wys[ed[v][i]] + 1);
	}
	oj[v] = 0;
}

void DFSA(int v, int dis)
{
	oj[v] = 1;
	int m1 = 0, m2 = 0;
	for(int i = 0; i < (int)ed[v].size(); ++i)
	{
		if(oj[ed[v][i]]) continue;
		int a = wys[ed[v][i]] + 1;
		if(a > m1) swap(m1, a);
		if(a > m2) swap(m2, a);
	}
	//cout << "DFSA: " << v << " " << dis << " | " << m1 << " " << m2 << "\n";
	

	ansc[v] = 0;
	if(dis - 1 > m1)
		ansc[v] = Query(1, 0, N - 1, 1, dis - 1 - m1);
	//cout << ansc[v] << " " << Query(1, 0, N - 1, 1, dis - 1) << "\n";
	for(int i = 0; i < (int)ed[v].size(); ++i)
	{
		if(oj[ed[v][i]]) continue;
		int l = m1;
		if(wys[ed[v][i]] + 1 == m1) l = m2;

		if(l > 0 && dis > 1)
			Change(1, 0, N - 1, max(1, dis - l), dis - 1, 1);

		int prv = lst[tab[v]];
		bool xd = 0;
		if(lst[tab[v]] == 0) xd = 1;
		if(!xd) xd = Check(lst[tab[v]]);
		if(xd)
		{
			lst[tab[v]] = dis;
			Add(dis, 1);
		}
		//cout << v << "->" << ed[v][i] << " " << xd << "\n";

		DFSA(ed[v][i], dis + 1);

		if(l > 0 && dis > 1)
			Change(1, 0, N - 1, max(1, dis - l), dis - 1, -1);
		if(xd)
		{
			lst[tab[v]] = prv;
			Add(dis, -1);
		}
	}


	oj[v] = 0;
}

void Solve()
{
    int n, m, a, b;
    cin >> n >> m;
    for(int i = 1; i < n; ++i)
    {
    	cin >> a >> b;
    	ed[a].pb(b);
    	ed[b].pb(a);
    }
    for(int i = 1; i <= n; ++i)
    	cin >> tab[i];
    oj[1] = 1;
    DFS(1);
    //cout << "E: " << sa << " " << sb << "\n";
    BFS(n, sa, sb);


    //cout << "FA: " << sa << "\n";
    DFSW(sa);
    DFSA(sa, 1);
    for(int i = 1; i <= n; ++i)
    	if(clo[i] == 1)
    		answer[i] = ansc[i];
    for(int i = 1; i <= m; ++i) lst[i] = 0;

    //cout << "\nFB: " << sb << "\n";
    DFSW(sb);
    DFSA(sb, 1);
    for(int i = 1; i <= n; ++i)
    	if(clo[i] == 2)
    		answer[i] = ansc[i];


    for(int i = 1; i <= n; ++i)
    	cout << answer[i] << "\n";
    //cout << "\n";
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    //int t; cin >> t;
    //while(t--)
        Solve();

    return 0;
}

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