Untitled
unknown
plain_text
a year ago
3.0 kB
8
Indexable
#include <bits/stdc++.h>
using namespace std;
#define ll long long
const int N = 2e5 + 5;
vector<int> adj[N];
struct LCA {
int LOG{};
vector<vector<int>> up;
vector<int> depth;
LCA(int n , int root) {
while ((1 << LOG) <= n) LOG++;
up = vector<vector<int>>(n, vector<int>(LOG));
depth = vector<int>(n);
dfs(root , 0);
}
void dfs(int root , int p) {
for (auto i: adj[root]) {
if(i == p)continue;
depth[i] = depth[root] + 1;
up[i][0] = root;
for (int j = 1; j < LOG; ++j) {
up[i][j] = up[up[i][j - 1]][j - 1];
}
dfs(i , root);
}
}
int lca(int a, int b) {
if (depth[a] < depth[b])swap(a, b);
int dif = depth[a] - depth[b];
for (int i = LOG - 1; i >= 0; --i) {
if (dif & (1 << i))
a = up[a][i];
}
if (a == b)return a;
for (int i = LOG - 1; i >= 0; --i) {
if (up[a][i] != up[b][i])
a = up[a][i],
b = up[b][i];
}
return up[a][0];
}
};
int timer{0};
int n , q , tin[N], tout[N] , c[N];
vector<int>fin[N] , fout[N];
void dfs(int u, int p) {
fin[c[u]].push_back(timer);
tin[u] = timer++;
for (auto v: adj[u])
if (v != p)
dfs(v, u);
fout[c[u]].push_back(timer++);
tout[u] = timer;
}
ll calc(bool ok , int u , int v, int val){
if(ok)
return std::upper_bound(fin[val].begin(), fin[val].end(), u) - std::lower_bound(fin[val].begin(), fin[val].end(), v);
return std::upper_bound(fout[val].begin(), fout[val].end(), u) - std::lower_bound(fout[val].begin(), fout[val].end(), v);
}
void solve() {
cin >> n >> q;
for (int i = 0; i < n; ++i) {
cin >> c[i];
c[i]--;
}
for (int i = 0; i < n - 1; ++i) {
int u ,v;
cin >> u >> v;
u-- , v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
dfs( 0 , 0);
LCA lca(n , 0);
for (int i = 0; i < q; ++i) {
int u , v , val;
cin >> u >> v >> val;
--u , --v , --val;
int LC = tin[lca.lca(u, v)];
// lca to u
int ti = tin[u];
int f = calc(1 , ti , LC , val) , s = calc(0 , ti , LC , val);
int ti2 = tin[v];
int f2 = calc(1 , ti2 , LC , val) , s2 = calc(0 , ti2 , LC , val);
// cerr << f << " " << s << " "<< f2 << " " << s2 << "\n";
f = f - s;
f2 = f2 - s2;
if(f + f2 > 0)
cout << 1;
else
cout << 0;
}
}
int main() {
ios::sync_with_stdio(false),
cin.tie(nullptr),
cout.tie(nullptr);
freopen("milkvisits.in", "r", stdin);
freopen("milkvisits.out", "w", stdout);
solve();
return 0;
}
Editor is loading...
Leave a Comment