#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;
}
