K Path Query - Codechef Solution | AskTheCode

Codechef July Long Challenge 2021 Solution | K Path Query solution in C++ | AskTheCode

---

Problem:

You're given a tree with N vertices numbered from 1 to N. Your goal is to handle Q queries. For each query, you are given K nodes v1,v2,…,vK. Find if there exists a simple path in the tree covering all the given vertices.

---

Input:

• The first line contains a single integer T - the number of test cases. The description of T test cases follows.

• The first line of each test case contains a single integer N.

• Each of the following N−1lines contains two integers u and v denoting an edge between nodes u and v.

• The next line contains a single integer Q - the number of queries.

• The next Q lines describe queries. The i-th line describes the i-th query and starts with the integer Ki — the number of vertices in the current query. Then Ki integers follow: v1,v2,…,vKi.

Output:

For each query print "YES" (without quotes) if there exists a simple path in the tree covering all the given nodes, otherwise print "NO" (without quotes).

You may print each character of the string in uppercase or lowercase (for example, the strings "yEs", "yes", "Yes" and "YES" will all be treated as identical).

---

Sample Input:

1
6
1 2
1 3
2 4
2 5
3 6
2
3 6 2 4
4 2 3 4 5

Sample Output:

YES
NO

---

EXPLANATION:

The structure of the given tree is shown below.

• For the first query, the path will be 4−2−1−3−6.

• For the second query, there is no simple path that covers all the given vertices.

---

Code:

#include<iostream>
#include<vector>
#include<utility>
#include<algorithm>
using namespace std;
const int N = 100009;
vector<int> g[N];
int parent[N], level[N], binaryLift[N][20]; // say logN = 20
int timeIn[N], timeOut[N]; // to calculate [is_ancestor]
int n, a, b, timer;
void dfs(int s, int par = -1) {
	parent[s] = par;
	timeIn[s] = timer++;
	if (par == -1) {
		binaryLift[s][0] = s;
		level[s] = 0;
	}
	else {
		binaryLift[s][0] = par;
		level[s] = level[par] + 1;
	}
	for (auto nx : g[s]) {
		if (nx == par)continue;
		dfs(nx, s);
	}
	timeOut[s] = timer++;
}
void create_ances_table() {
	for (int lg = 1; lg < 20; lg++) {
		for (int i = 0; i < n; i++) {
			binaryLift[i][lg] = binaryLift[binaryLift[i][lg - 1]][lg - 1];
		}
	}
}
bool is_ancestor(int a, int b) {
	if (timeIn[a] <= timeIn[b] && timeOut[b] <= timeOut[a]) return true;
	return false;
}
int lca(int a, int b) {
	if (is_ancestor(a, b))return a;
	if (is_ancestor(b, a))return b;
	for (int lg = 19; lg >= 0; lg--) {
		if (!is_ancestor(binaryLift[a][lg], b)) {
			a = binaryLift[a][lg];
		}
	}
	return binaryLift[a][0];
}
bool cmp(int a, int b) {
	return level[a] < level[b];
}
vector<int> nodes;
vector<int> path1, path2;
int main() {
	int testcase;
	cin >> testcase;
	for (int test = 1; test <= testcase; test++) {
		cin >> n;
		for (int i = 0; i < n; i++) {
			g[i].clear();
		}
		timer = 0;
		for (int i = 1; i < n; i++) {
			cin >> a >> b;
			a--;
			b--;
			g[a].push_back(b);
			g[b].push_back(a);
		}
		dfs(0);
		create_ances_table();
		int q;
		cin >> q;
		while (q--) {
			int k;
			cin >> k;
			nodes.clear();
			path1.clear();
			path2.clear();
			while (k--) {
				cin >> a;
				a--;
				nodes.push_back(a);
			}
			sort(nodes.begin(), nodes.end(), cmp);
			int u = nodes.back();
			nodes.pop_back();
			path1.push_back(u);
			while (nodes.size()) {
				int ances = nodes.back();
				nodes.pop_back();
				if (is_ancestor(ances, u)) {
					path1.push_back(ances);
					u = ances;
				}
				else {
					path2.push_back(ances);
				}
			}
			bool ok = true;
			for (int i = 1; i < path1.size() && ok; i++) {
				if (level[path1[i - 1]] == level[path1[i]])ok = false;
				if (!is_ancestor(path1[i], path1[i - 1]))ok = false;
			}
			for (int i = 1; i < path2.size() && ok; i++) {
				if (level[path2[i - 1]] == level[path2[i]])ok = false;
				if (!is_ancestor(path2[i], path2[i - 1]))ok = false;
			}
			if (path1.size() && path2.size()) {
				if(lca(path1[0],path2[0]) != lca(path1[0],path2[path2.size()-1])) ok = false;
				if (lca(path2[0], path1[0]) != lca(path2[0], path1[path1.size() - 1])) ok = false;
			}
			if (ok) {
				cout << "YES\n";
			}
			else {
				cout << "NO\n";
			}
		}
	}
}