Codeforces Round #620 (Div. 2) E. 1-Trees and Queries

출처 : https://codeforces.com/contest/1304/problem/E

LCA를 이용하면 댄다...

풀이는

int query(int nu, int nv, int u, int v, int len)

함수에 주석으로 적어놓았다....

아래는 전체 코드

#include<bits/stdc++.h>
#define endl "\n"
using namespace std;
using ll = long long;
struct TREE
{
	int vn;
	vector<vector<int>>elist;
	vector<int>s, p, cn, cd, ci, cc, d;
	TREE(int a) : vn(a)
	{
		elist.resize(vn + 1);
		s.resize(vn + 1, 0);
		p.resize(vn + 1, 0);
		cn.resize(vn + 1, 0);
		cd.resize(vn + 1, 0);
		ci.resize(vn + 1, 0);
		cc.resize(vn + 1, 0);
		d.resize(vn + 1, 0);
	}
	void pb(int a, int b)
	{
		elist[a].push_back(b);
	}
	void build_hld()
	{
		go1(1, 0);
		go2(1, 0, 1, 0);
	}
	int go1(int c, int pa)
	{
		s[c] = 1;
		p[c] = pa;
		for (auto& nx : elist[c])
		{
			if (nx != pa)
			{
				d[nx] = d[c] + 1;
				s[c] += go1(nx, c);
			}
		}
		return s[c];
	}
	void go2(int c, int pa, int ccn, int ccd)
	{
		int h = -1;
		cn[c] = ccn;
		cd[c] = ccd;
		ci[c] = cc[ccn]++;
		for (auto& nx : elist[c])
		{
			if (nx != pa && (h == -1 || s[nx] > s[h]))
			{
				h = nx;
			}
		}
		if (h != -1)
		{
			go2(h, c, ccn, ccd);
		}
		for (auto& nx : elist[c])
		{
			if (nx != pa && nx != h)
			{
				go2(nx, c, nx, ccd + 1);
			}
		}
	}
	int lca(int a, int b)
	{
		while (cn[a] != cn[b])
		{
			if (cd[a] > cd[b])
			{
				a = p[cn[a]];
			}
			else
			{
				b = p[cn[b]];
			}
		}
		return ci[a] > ci[b] ? b : a;
	}
	int dis(int a, int b)
	{
		return d[a] + d[b] - (d[lca(a, b)] << 1);
	}
};
struct SOLUTION
{
	int vn;
	TREE t;
	SOLUTION(int a) : vn(a), t(a)
	{
		input();
		t.build_hld();
		proc();
	}
	void input()
	{
		int i, a, b;
		for (i = 0; i < vn - 1; i++)
		{
			cin >> a >> b;
			t.pb(a, b);
			t.pb(b, a);
		}
	}
	void proc()
	{
		int qn, x, y, a, b, k;
		cin >> qn;
		while (qn--)
		{
			cin >> x >> y >> a >> b >> k;
			if (query(x, y, a, b, k))
			{
				cout << "YES" << endl;
			}
			else
			{
				cout << "NO" << endl;
			}
		}
	}
	int query(int nu, int nv, int u, int v, int len)
	{
		int dis;
		//origin path u - > v
		dis = t.dis(u, v);
		if (dis == len || (dis < len && !((len - dis) & 1)))
		{
			return 1;
		}

		// dis(nu, nv) == dis(nv, nu) == 1;

		//new path1 : u - > nu - > nv - > v
		dis = t.dis(u, nu) + 1 + t.dis(nv, v);
		if (dis == len || (dis < len && !((len - dis) & 1)))
		{
			return 1;
		}
		//new path2 : u - > nv - > nu - > v
		dis = t.dis(u, nv) + 1 + t.dis(nu, v);
		if (dis == len || (dis < len && !((len - dis) & 1)))
		{
			return 1;
		}

		/*
		if (dis == len || (dis < len && !((len - dis) & 1)))
		{
			return 1;
		}
		u - > v 거리가 len 이면 도착 성공

		or

		u -> v 까지 최단으로 갔을때 더 움직여야할 거리가 남았고 (dis < len) 그 남은 거리가 짝수이면 성공( !((len - dis) & 1)) )
		왜냐면 두 노드사이를 왔다갔다 반복하면 된다.
		
		아니면? 안된다 return 0;
		*/
		return 0;
	}
};
void make_answer();
int main(void)
{
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	make_answer();
	return 0;
}
void make_answer()
{
	int a;
	cin >> a;
	SOLUTION sol(a);
}