洛谷P5838 [USACO19DEC]Milk Visits G 题解

题意：

Farmer John 计划建造 \(N\) 个农场，用 \(N-1\) 条道路连接，构成一棵树（也就是说，所有农场之间都互相可以到达，并且没有环）。每个农场有一头奶牛，品种为 \(1\) 到 \(N\) 之间的一个整数 \(T_i\) 。

Farmer John 的 \(M\) 个朋友经常前来拜访他。在朋友 \(i\) 拜访之时，Farmer John 会与他的朋友沿着从农场 \(A_i\) 到农场 \(B_i\) 之间的唯一路径行走（可能有 \(A_i = B_i\)）。除此之外，他们还可以品尝他们经过的路径上任意一头奶牛的牛奶。由于 Farmer John 的朋友们大多数也是农场主，他们对牛奶有着极强的偏好。他的每个朋友都只喝某种特定品种的奶牛的牛奶。任何 Farmer John 的朋友只有在他们访问时能喝到他们偏好的牛奶才会高兴。

请求出每个朋友在拜访过后是否会高兴。

输入格式：

输入的第一行包含两个整数 \(N\) 和 \(M\)。

第二行包含 \(N\) 个空格分隔的整数 \(T_1,T_2,\ldots, T_N\)。第 \(i\) 个农场内的奶牛的品种用 \(T_i\) 表示。

以下 \(N-1\) 行，每行包含两个不同的整数 \(X\) 和 \(Y\)（\(1 \leq X, Y \leq N\)），表示农场 \(X\) 与 \(Y\) 之间有一条边。

以下 \(M\) 行，每行包含整数 \(A_i\)，\(B_i\)，以及\(C_i\)。\(A_i\) 和 \(B_i\) 表示朋友 \(i\) 拜访时行走的路径的端点，\(C_i\)（\(1\le C_i\le N\)）表示这个朋友喜欢的牛奶的奶牛品种。

输出格式：

输出一个长为 \(M\) 的二进制字符串。如果第 \(i\) 个朋友会感到高兴，则字符串的第 \(i\) 个字符为 1，否则为 0。

数据范围：

\(1 \leq N \leq 10^5\)，\(1 \leq M \leq 10^5\)。

知道为啥我刚要写个 Tarjan 求 LCA 了吧，就是因为这题的最优解用到了它的思想。

注意到 \(u,v\) 间有目标颜色，记 \(d=\mathrm{LCA}(u,v)\) ，则 \([u,d]\cup[d,v]\) 内至少有一个。

考虑维护 \(1 \leadsto u\) 中每个颜色最深在哪里，记它在 \(1 \leadsto u\) 上的下个节点为 \(\mathtt{top}\) 。

考虑当询问的左节点被访问时记录该颜色对应的 \(\mathtt{top}\) ，然后当右节点被访问时校验是否相等

容易发现，如果 \(\mathtt{top}\) 不相等，则这个点在 \(u \leadsto v\) 上（记录「下个节点」的原因，考虑 \(\mathrm{LCA}(u,v)\) ）

时间复杂度 \(\mathcal{O}(n + m)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
#define N ((int)(1e5 + 15))

vector<int> vec[N];
int n,m,pos=1,head[N],c[N],q[N],top[N],ans[N];
struct Edge { int u,v,next; } e[N * 2];
void addEdge(int u,int v) { e[++pos] = {u,v,head[u]}; head[u] = pos; }
void dfs(int u,int f)
{
    int tmp = top[c[u]];
    for(int i : vec[u])
    {
        if(ans[i] != -1) ans[i] = (ans[i] != top[q[i]]);
        else ans[i] = top[q[i]];
    }
    for(int i = head[u]; i; i = e[i].next)
    { int v = e[i].v; if(v != f) { top[c[u]] = v; dfs(v,u); } }
    top[c[u]] = tmp;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n >> m;
    for(int i = 1; i <= n; i++) cin >> c[i];
    for(int i = 1, u,v; i < n; i++) { cin >> u >> v; addEdge(u,v); addEdge(v,u); }
    for(int i = 1,u,v; i <= m; i++)
    {
        cin >> u >> v >> q[i];
        if(c[u] == q[i] || c[v] == q[i]) ans[i] = 1;
        else { ans[i] = -1, vec[u].emplace_back(i); vec[v].emplace_back(i); }
    }
    dfs(1,0); for(int i = 1; i <= m; i++) cout << ans[i];
    return 0;
}

当然这道题也是有「通俗易懂的数据结构解法」的啦。

题目中的「询问 \(u \leadsto v\) 是否存在颜色为 \(c\) 的节点」，其实就是序列上 \([l,r]\) 是否存在颜色为 \(c\) 的位置。

这种我们直接用 vector 维护每个颜色的出现位置，然后二分一下就好了，这道题无非就是套了个树剖。

时间复杂度 \(\mathcal{O}(n \log^2 n)\) ，除非特意构造否则一般难卡满。

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
#define N ((int)(1e5 + 15))

vector<int> vec[N];
int n,m,pos=1,head[N],c[N],top[N],id[N],fa[N],son[N],sz[N],dep[N];
struct Edge { int u,v,next; } e[N * 2];
void addEdge(int u,int v) { e[++pos] = {u,v,head[u]}; head[u] = pos; }
void dfs(int u,int f,int d)
{
    dep[u] = d; fa[u] = f; sz[u] = 1; int mx = -1;
    for(int i = head[u]; i; i = e[i].next)
    {
        int v = e[i].v; if(v == f) continue;
        dfs(v,u,d+1); sz[u] += sz[v]; if(sz[v] > mx) { mx = sz[v], son[u] = v; }
    }
}
void dfs(int u,int ftop)
{
    static int tot = 0; top[u] = ftop; id[u] = ++tot;
    vec[c[u]].emplace_back(tot); if(son[u]) dfs(son[u], ftop);
    for(int i = head[u]; i; i = e[i].next)
    { int v = e[i].v; if(v != fa[u] && v != son[u]) dfs(v,v); }
}
bool query(int nl,int nr,int k)
{ return *lower_bound(vec[k].begin(), vec[k].end(), nl) <= nr; }
bool qRange(int x,int y,int k)
{
    while(top[x] != top[y])
    {
        if(dep[top[x]] < dep[top[y]]) swap(x,y);
        if(query(id[top[x]], id[x], k)) return 1; x = fa[top[x]];
    }
    if(dep[x] > dep[y]) swap(x,y);
    return query(id[x], id[y], k);
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n >> m;
    for(int i = 1; i <= n; i++) cin >> c[i];
    for(int i = 1, u,v; i < n; i++)
    { cin >> u >> v; addEdge(u,v); addEdge(v,u); }
    dfs(1,0,1); dfs(1,1); for(int i = 1; i <= n; i++) vec[i].emplace_back(INF);
    for(int i = 1, u,v,w; i <= m; i++)
    { cin >> u >> v >> w; cout << qRange(u,v,w); }
    return 0;
}