#include "werewolf.h"
#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for (int i = a; i < b; i++)
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define srt(x) sort(all(x))
typedef long long ll;
using vi = vector<int>;
using vl = vector<ll>;
/*
A persistent DSU implementation that's slow and bad (copies the DSU each time)
*/
struct DSU {
int n;
vector<int> par, size;
vector<bitset<3000>> has;
DSU(int N) {
n = N;
par.assign(n, 0); iota(par.begin(), par.end(), 0);
size.assign(n, 1);
has.assign(n, bitset<3000>());
FOR(i, 0, n) has[i].flip(i);
}
int find_set(int v) {
if (par[v] == v) {
return v;
}
return find_set(par[v]);
}
void unite(int a, int b) {
a = find_set(a);
b = find_set(b);
if (a != b) {
if (size[a] < size[b]) swap(a, b);
par[b] = a;
size[a] += size[b];
has[a] |= has[b];
}
}
};
struct PersistentDSU {
int n, timer;
vector<DSU> states;
PersistentDSU(int N) {
n = N;
timer = 0;
states.push_back(DSU(n));
}
int find_set(int v, int time) {
return states[time].find_set(v);
}
void unite(int a, int b) {
DSU nDSU = states.back();
nDSU.unite(a, b);
states.push_back(nDSU);
timer++;
}
bool does_contain(int head, int node, int time) {
return states[time].has[head].test(node);
}
};
/*
N = num nodes
X, Y = edges
S, E = start and end nodes
L, R = human and werewolf cities
*/
/*
Solution:
- for a given value of l, each node u has a connected set of nodes H_u they can reach as a human (<= R)
- for a given value of r, each node u has a connected set of nodes V_u they can reach as a vampire (>= L)
- a query (start, end, L, R) is possible iff |H_u ∩ V_u| ≥ 1
Implementation (slow):
(1) persistent dsu data structure which supports set intersection queries
(2) we just need 2 persistent dsus, one where we add nodes from 0 to n (red, vampire), and the other from n to 0 (blue, human), then do queries off
(3) to answer a query (start, end, left, right) we need to find the set intersection of
blue_dsu[time = left][head of start].members and the set red_dsu[time = right[head of end] and count the number of intersections
*/
int n, m, q;
vi x, y, s, e, l, r;
vi queries;
vi check_validity(int N, vi X, vi Y, vi S, vi E, vi L, vi R) {
n = N; x = X; y = Y; s = S; e = E; l = L; r = R;
m = x.size(), q = S.size();
vi A(q);
auto blueDSU = PersistentDSU(n);
auto redDSU = PersistentDSU(n);
vi wB(m), wR(m); iota(all(wB), 0); iota(all(wR), 0);
sort(all(wR), [&](int a, int b) {return max(X[a], Y[a]) < max(X[b], Y[b]);});
sort(all(wB), [&](int a, int b) {return min(X[a], Y[a]) > min(X[b], Y[b]);});
vi maxR(m), minB(m);
FOR(i, 0, m) {
maxR[i] = max(X[wR[i]], Y[wR[i]]);
minB[i] = min(X[wB[i]], Y[wB[i]]);
}
for (auto &edge : wR) {
redDSU.unite(X[edge], Y[edge]);
}
for (auto &edge : wB) {
blueDSU.unite(X[edge], Y[edge]);
}
FOR(query, 0, q) {
if (S[query] < L[query] || E[query] > R[query]) {
A[query] = 0;
continue;
}
// int timeRed = 0; // do binary search
// int timeBlue = 0; // do binary search
// timeRed = # of red-edges with max ≤ R[i]
int timeRed = int(upper_bound(maxR.begin(), maxR.end(), R[query]) - maxR.begin());
// timeBlue = # of blue-edges with min ≥ L[i]
int lo = 0, hi = m;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (minB[mid] >= L[query]) lo = mid + 1;
else hi = mid;
}
int timeBlue = lo;
int red_parent = redDSU.find_set(E[query], timeRed);
int blue_parent = blueDSU.find_set(S[query], timeBlue);
bool can = (redDSU.states[timeRed].has[red_parent] & blueDSU.states[timeBlue].has[blue_parent]).any();
A[query] = can ? 1 : 0;
}
return A;
}
