Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- using namespace std;
- struct CentroidDecomposer {
- int n;
- vector<set<int>> G;
- vector<int> Size;
- // Depth and parent in the centroid tree
- vector<int> CD, CP;
- vector<vector<int>> dp, depth;
- vector<bool> tag;
- vector<int> all;
- vector<int> down, up;
- int k;
- CentroidDecomposer(int n, int k) :
- n(n), k(k), G(n), Size(n), CD(n), CP(n), tag(n, false),
- down(n, -1), up(n, -1),
- depth(20, vector<int>(n)), dp(20, vector<int>(n)) {}
- void Tag(int node) {
- tag[node] = true;
- }
- void AddEdge(int a, int b) {
- G[a].insert(b);
- G[b].insert(a);
- }
- bool Solve(int a, int b) {
- int na = a, nb = b;
- while (CD[na] > CD[nb])
- na = CP[na];
- while (CD[na] < CD[nb])
- nb = CP[nb];
- while (na != nb) {
- na = CP[na];
- nb = CP[nb];
- }
- int cen = na;
- int da = dp[CD[cen]][a];
- int db = dp[CD[cen]][b];
- if (da + db <= k) {
- return true;
- }
- return false;
- }
- int rec_size, rec_centr;
- void ComputeSizeAndCentroid(int node, int par) {
- Size[node] = 1;
- int max_size = 0;
- for(auto vec : G[node])
- if(vec != par) {
- ComputeSizeAndCentroid(vec, node);
- Size[node] += Size[vec];
- max_size = max(max_size, Size[vec]);
- }
- max_size = max(max_size, rec_size - Size[node]);
- if(2 * max_size <= rec_size)
- rec_centr = node;
- }
- void DoUsefulDFS(int node, int par, int cd) {
- Size[node] = 1;
- all.push_back(node);
- for(auto vec : G[node])
- if(vec != par) {
- depth[cd][vec] = 1 + depth[cd][node];
- DoUsefulDFS(vec, node, cd);
- Size[node] += Size[vec];
- }
- }
- void DFS1(int cd, int node, int par) {
- if (tag[node]) {
- down[node] = node;
- }
- for (auto vec : G[node]) {
- if (vec == par) continue;
- DFS1(cd, vec, node);
- if (down[vec] != -1 && (down[node] == -1 ||
- depth[cd][down[node]] > depth[cd][down[vec]])) {
- down[node] = down[vec];
- }
- }
- }
- multimap<int, int> DFS2(int cd, int node, int par) {
- multimap<int, int> big;
- big.insert({depth[cd][node], node});
- for (auto vec : G[node]) {
- if (vec == par) continue;
- auto small = DFS2(cd, vec, node);
- if (small.size() > big.size()) swap(small, big);
- for (auto x : small) {
- big.insert(x);
- }
- }
- if (down[node] != -1) {
- int dd = depth[cd][down[node]];
- int dr = depth[cd][node];
- // dd + x - 2 * dr <= k
- // x <= k + 2 * dr - dd
- int maxd = k + 2 * dr - dd;
- int mind = dd + 1;
- auto it = big.lower_bound(mind);
- while (it != big.end() && it->first <= maxd) {
- assert(up[it->second] == -1);
- up[it->second] = down[node];
- it = big.erase(it);
- }
- }
- return big;
- }
- void BFS(int cd, int root) {
- DFS1(cd, root, -1);
- DFS2(cd, root, -1);
- sort(all.begin(), all.end(), [&](int a, int b) {
- return depth[cd][a] < depth[cd][b];
- });
- for (auto x : all) {
- dp[cd][x] = depth[cd][x];
- if (up[x] != -1) {
- assert(depth[cd][up[x]] < depth[cd][x]);
- assert(dp[cd][up[x]] != -1);
- dp[cd][x] = dp[cd][up[x]];
- }
- }
- // cerr << "CENTROID: " << root + 1 << endl;
- // cerr << "SPEC: ";
- // for (auto x : spec) {
- // cerr << x + 1 << " ";
- // }
- // cerr << endl;
- // for (auto x : all) {
- // cerr << x + 1 << ": " << endl;
- // cerr << "down: " << down[x] + 1 << endl;
- // cerr << "up: " << up[x] + 1 << endl;
- // cerr << "dp: " << dp[cd][x] << "\n";
- // // cerr << "bfsr: " << bfsr[x] << '\n';
- // // cerr << "bfsd: " << bfsd[x] << '\n';
- // }
- // cerr << endl;
- for (auto x : all) {
- down[x] = -1;
- up[x] = -1;
- }
- }
- void RecDecomp(int node, int size, int cp, int cd) {
- // node -> centroid(node)
- rec_size = size;
- ComputeSizeAndCentroid(node, -1);
- node = rec_centr;
- CP[node] = cp;
- CD[node] = cd;
- // You can do whichever DFS you want here
- all.clear();
- DoUsefulDFS(node, -1, cd);
- BFS(cd, node);
- for(auto vec : G[node]) {
- G[vec].erase(node);
- RecDecomp(vec, Size[vec], node, cd + 1);
- }
- }
- void Decompose(int root = 0) {
- RecDecomp(root, n, -1, 0);
- for(auto x : CD)
- assert(x < 20);
- }
- };
- int main() {
- ios_base::sync_with_stdio(false);
- cin.tie(0);
- int n, k, r; cin >> n >> k >> r;
- CentroidDecomposer cd(n, k);
- for (int i = 1; i < n; ++i) {
- int a, b; cin >> a >> b;
- cd.AddEdge(a - 1, b - 1);
- }
- for (int i = 0; i < r; ++i) {
- int x; cin >> x; cd.Tag(x - 1);
- }
- cd.Decompose();
- int q; cin >> q;
- while (q--) {
- int a, b; cin >> a >> b;
- cout << (cd.Solve(a - 1, b - 1) ? "YES" : "NO") << '\n';
- }
- return 0;
- }
Add Comment
Please, Sign In to add comment