#include "werewolf.h" #include <algorithm> #include <functional> #include <numeric> std::vector<int> check_validity(int N, std::vector<int> X, std::vector<int> Y, std::vector<int> S, std::vector<int> E, std::vector<int> L, std::vector<int> R) { // construct graph std::vector<std::vector<int>> adj(N + 1); int M = X.size(); for (int i = 0; i < M; ++i) { adj[X[i]].push_back(Y[i]); adj[Y[i]].push_back(X[i]); } // Construct KRTs std::vector<int> par(N + 1); std::vector<int> sz(N + 1); std::function<int(int)> find_set; find_set = [&](int v) -> int { if (v == par[v]) return v; return par[v] = find_set(par[v]); }; int dfsTimer; // construct KRT for left threshold std::iota(par.begin(), par.end(), 0); std::fill(sz.begin(), sz.end(), 0); std::vector<std::vector<int>> tree1(N + 1); for (int i = N; i >= 1; --i) { sz[i] = 1; for (int j : adj[i]) if (sz[j]) { int p = find_set(j); tree1[i].push_back(p); } for (int j : adj[i]) if (sz[j]) { int p = find_set(j); par[p] = i; sz[i] += sz[p]; } } dfsTimer = 0; std::vector<int> tin1(N + 1), tout1(N + 1), A(N + 1); std::vector<int> posA(N + 1); std::function<void(int)> DFS1; DFS1 = [&](int u) { tin1[u] = ++dfsTimer; A[tin1[u]] = u; posA[u] = tin1[u]; for (int v : tree1[u]) DFS1(v); tout1[u] = dfsTimer; }; DFS1(1); // construct KRT for right threshold std::iota(par.begin(), par.end(), 0); std::fill(sz.begin(), sz.end(), 0); std::vector<std::vector<int>> tree2(N + 1); for (int i = 1; i <= N; ++i) { sz[i] = 1; for (int j : adj[i]) if (sz[j]) { int p = find_set(j); tree2[i].push_back(p); } for (int j : adj[i]) if (sz[j]) { int p = find_set(j); par[p] = i; sz[i] += sz[p]; } } dfsTimer = 0; std::vector<int> tin2(N + 1), tout2(N + 1), B(N + 1); std::function<void(int)> DFS2; DFS2 = [&](int u) { tin2[u] = ++dfsTimer; B[tin2[u]] = u; for (int v : tree2[u]) DFS1(v); tout2[u] = dfsTimer; }; DFS2(N); // Sweep line to find answer int Q = S.size(); std::vector<std::pair<int,int>> queries(Q); for (int i = 0; i < Q; ++i) queries[i] = {S[i], E[i]}; std::vector<int> BIT(N + 1, 0); auto update = [&](int pos) -> void { for (; pos <= N; pos += pos & (-pos)) BIT[pos]++; }; auto sum = [&](int pos) -> int { int ret = 0; for (; pos > 0; pos -= pos & (-pos)) ret += BIT[pos]; return ret; }; auto query = [&](int l, int r) -> int { if (l > r) return 0; return sum(r) - sum(l - 1); }; std::vector<std::vector<int>> queL(N + 1); for (int i = 0; i < Q; ++i) queL[tin2[E[i]] - 1].push_back(i); std::vector<std::vector<int>> queR(N + 1); for (int i = 0; i < Q; ++i) queR[tout2[E[i]]].push_back(i); std::vector<int> answer(Q, 0); for (int i = 1; i <= N; ++i) { update(posA[B[i]]); for (int qID : queL[i]) { int u = queries[qID].first; answer[qID] -= query(tin1[u], tout1[u]); } for (int qID : queR[i]) { int u = queries[qID].first; answer[qID] += query(tin1[u], tout1[u]); } } return answer; }