Luogu P5838 [USACO19DEC]Milk Visits G

题目描述

https://www.luogu.com.cn/problem/P5838

简要题意：给定一棵 $n$ 个节点的无根树，每个节点有一个颜色 $c_i$，现在有 $m$ 次询问，每次询问给定 $x,y,z$，问从 $x$ 到 $y$ 的简单路径上是否有颜色 $z$

$n,m\le 10^5$

Solution

我们考虑在 $dfs$ 的过程中维护当前节点到根的路径上每种颜色的深度最深的出现位置，那么对于一次询问，如果对于 $x$ 节点到根路径上 $z$ 出现的位置与 $y$ 到根路径上 $z$ 出现的位置相同，那么除非这个位置是 $x$ 和 $y$ 的 $lca$，否则说明 $z$ 没有出现在 $x$ 到 $y$ 的简单路径上，出现位置不同则说明 $z$ 一定出现在 $x$ 到 $y$ 的简单路径上

但是求 $lca$ 太麻烦了，所以我们在 $dfs$ 的时候把颜色放到儿子节点上，这样就不用特判 $lca$ 了，时间复杂度 $O(n)$

#include <iostream>
#include <vector>
#define maxn 100010
using namespace std;

int n, m, c[maxn];

struct Edge {
    int to, next;
} e[maxn * 2]; int c1, head[maxn];
void add_edge(int u, int v) {
    e[c1].to = v; e[c1].next = head[u]; head[u] = c1++;
}

vector<pair<int, int>> A[maxn]; int ans[maxn];

int top[maxn];
void dfs(int u, int fa) {
    int _top = top[c[u]];
    for (int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to; if (v == fa) continue;
        top[c[u]] = v; dfs(v, u); 
    } top[c[u]] = _top;
    for (auto [col, id] : A[u]) {
        if (ans[id] == -1) ans[id] = top[col];
        else ans[id] = ans[id] != top[col];
    }
}
 
int main() { fill(head, head + maxn, -1); 
    ios::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);

    cin >> n >> m;
    for (int i = 1; i <= n; ++i) cin >> c[i];
    for (int i = 1; i < n; ++i) {
        int x, y; cin >> x >> y;
        add_edge(x, y); add_edge(y, x); 
    } 
    for (int i = 1; i <= m; ++i) {
        int x, y, z; cin >> x >> y >> z; 
        if (c[x] == z || c[y] == z) ans[i] = 1;
        else ans[i] = -1, A[x].emplace_back(z, i), A[y].emplace_back(z, i);
    } dfs(1, 0);
    for (int i = 1; i <= m; ++i) cout << ans[i]; cout << "\n";
    return 0; 
}