Milk Visits (USACO Gold 2019 December)

此题是比较灵活的一个LCA应用，问题可以表述为：

查询树上从 $a$ 到 $b$ （上图中o组成的路），是否存在目标颜色的点（ $c$ ）。

和这个题的silver版本类似，我们可以用离线查询，也就是不是按照查询顺序，而是按方便我们的顺序来解决，会让问题更好解一些。具体来说，就是按DFS的顺序来解。DFS到一个节点时，回答从这个点出发或者到达这个点的查询。

在DFS过程中，可以看到只有当前活跃路径（从根到 $a$ 或者 $b$ 的路径）上的节点才有关系，所以我们维护这条路径上的各种颜色的节点列表就行。实际上，只有深度最深的点才有用，所以对于每种颜色，我们用一个栈来维护这个列表。

用倍增找到LCA( $a$ , $b$ )之后，我们就是要判断栈顶点是否在 $a$ / $b$ 到LCA( $a$ , $b$ )这段路径上。这个的办法是 $O(1)$ 的，也就是直接比较深度就行了。

我的程序：

#include <bits/stdc++.h>
using namespace std;
 
// n, m <= 1e5
int n, m;
int T[100005];
vector<int> adj[100005];
struct F {
    int a,b,c;
};
F fs[100005];               // all friends
vector<int> v2f[100005];    // vertex to friend id
stack<int> xs[100005];      // 当前路径上每种颜色的点的位置（从上到下）
int ans[100005];            // 每个朋友的结果，是两个端点结果的or
 
struct LCA {
    int fa[100005][21];
    int depth[100005];
    void dfs(int s, int e, int d) {
        depth[s] = d;
        if (e >= 0) fa[s][0] = e;
        for (int u: adj[s]) 
            if (u == e) continue;
            else dfs(u, s, d+1);
    }
    void init() {                       
        dfs(0, -1, 0);                  // calc fa[x][0] and depth[]
        for (int i = 1; i <= 20; i++)   // build binary lifting pointers
            for (int j = 0; j < n; j++)
            fa[j][i] = fa[fa[j][i-1]][i-1];
    }
    int lca(int a, int b) {
        if (depth[a] < depth[b]) swap(a,b);
        for (int i = 20; i >= 0; i--)   // lift a to same depth of b
            if (depth[fa[a][i]] >= depth[b])
                a = fa[a][i];
        if (a == b) return a;
        for (int i = 20; i >= 0; i--)   // lift a and b together
            if (fa[a][i] != fa[b][i])
                a = fa[a][i], b = fa[b][i];
        return fa[a][0];
    }
};
LCA lca;
 
void dfs(int s, int e) {
    xs[T[s]].push(s);
    // 是f的端点时进行计算
    for (int i: v2f[s]) {
        int c = fs[i].c;
        if (xs[c].empty()) continue;
        int x = xs[c].top();
        if (lca.depth[x] >= lca.depth[lca.lca(fs[i].a, fs[i].b)])
            ans[i] = 1;
    }
 
    for (int u: adj[s]) {
        if (u == e) continue;
        dfs(u, s);
    }
    xs[T[s]].pop();
}
 
int main() {
    freopen("milkvisits.in", "r", stdin);
    freopen("milkvisits.out", "w", stdout);
    scanf("%d %d", &n, &m);
    for (int i = 0; i < n; i++)
        scanf("%d", &T[i]);
    for (int i = 0; i < n-1; i++) {
        int u,v;
        scanf("%d %d", &u, &v);
        u--,v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 0; i < m; i++) {
        scanf("%d %d %d", &fs[i].a, &fs[i].b, &fs[i].c);
        fs[i].a--; fs[i].b--;
        v2f[fs[i].a].push_back(i);
        v2f[fs[i].b].push_back(i);
    }
    lca.init();
    dfs(0, -1);
    for (int i = 0; i < m; i++)
        printf("%d", ans[i]);
 
    return 0;
}