洛谷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;
}