洛谷P5838 [USACO19DEC] Milk Visits G

洛谷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,\dots ,T_N$。第 $i$ 个农场内的奶牛的品种用 $T_i$ 表示。

以下 $N-1$ 行，每行包含两个不同的整数 $X$ 和 $Y$（$1\le X,Y\le 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

输入 #1

5 5
1 1 2 1 2
1 2
2 3
2 4
1 5
1 4 1
1 4 2
1 3 2
1 3 1
5 5 1

输出 #1

10110

输入输出样例 #2

输入 #2

6 4
1 2 3 3 3 3
1 2
2 3
3 4
2 5
5 6
4 6 1
4 6 2
4 6 3
4 6 4

输出 #2

0110

说明/提示

测试点性质：

测试点 $2$ 为以下第二个样例。

测试点 $3$ 满足 $N\le 10^3$，$M\le 2\cdot 10^3$。

测试点 $4\sim 7$ 满足 $C_i\le 10$。

对于 $100\%$ 的数据，$1\le N\le 10^5$，$1\le M\le 10^5$。

供题：Spencer Compton

思路详解

确定算法

首先，一眼看到$N$个点与$N-1$条边，很明显是树链剖分。
考虑如何求$(u,v)$的链上有无我们想要的颜色？？？
先想我们如果只有一个颜色，我们应如何判断一条链上有没有。显然，在不需要修改的情况下，我们只需要使用一个前缀和数组维护即可。
我们现在有多个颜色，那考虑对每个颜色开一个数组，可是我们发现若每个数组都开$10^5$的话，那我们的空间肯定会爆炸，那我们考虑使用动态数组，将每个节点$u$的$dfs$序$i{d}_u$压入对应颜色的动态数组即可。不过动态数组查询过于复杂，居然需要用到lower_bound这种高科技(doge)，所以我直接选用万能的$FHQTreap$。

具体实现

1. 对于$FHQTreap$：我们只需要维护一个有序序列，用于保存对应颜色节点的$dfs$序即可，还需要完成求在值域$[l,r]$内的数有几个。
2. 对于树剖：没什么好说的，就是普通的树剖模板，查询就是查询在$[u,v]$的链上有多少个要求颜色的个数即可。

code

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+5;
struct lit{int val,siz,pri;lit *l,*r;lit(int v){val=v;siz=1;pri=rand();l=r=nullptr;}
};
class FHQ{
private:lit *root;int getsize(lit *u){return u?u->siz:0;}void pushup(lit *&u){if(!u)return ;u->siz=1+getsize(u->l)+getsize(u->r);}void split(lit *u,int v,lit *&x,lit *&y){if(!u){x=y=nullptr;return ;}if(u->val<=v){x=u;split(u->r,v,x->r,y);pushup(x);}else{y=u;split(u->l,v,x,y->l);pushup(y);}}lit *merge(lit *x,lit *y){if(!x)return y;if(!y)return x;if(x->pri>y->pri){x->r=merge(x->r,y);pushup(x);return x;}else{y->l=merge(x,y->l);pushup(y);return y;}}void _print(lit *u){if(!u)return ;_print(u->l);cout<<u->val<<' ';_print(u->r);}
public:void print(){_print(root);cout<<'\n';}void insert(int v){lit *l,*r;split(root,v,l,r);root=merge(merge(l,new lit(v)),r);}int query_num(int x,int y){lit *l,*mid,*r;split(root,y,mid,r);split(mid,x-1,l,mid);int res=getsize(mid);root=merge(merge(l,mid),r);return res;}
}tr[N];//普通的FHQ Treap
vector<int>e[N];
int fa[N],siz[N],son[N],dep[N];
int id[N],reid[N],top[N],tim=0;
int co[N];
int n,m;
class HCD{//重链剖分
private:void dfs1(int u,int fath,int dp){dep[u]=dp;fa[u]=fath;siz[u]=1;for(int v:e[u]){if(v==fath)continue;dfs1(v,u,dp+1);siz[u]+=siz[v];if(siz[v]>siz[son[u]])son[u]=v;}}void dfs2(int u,int fath,int ori){id[u]=++tim;reid[tim]=u;top[u]=ori;if(son[u])dfs2(son[u],u,ori);for(int v:e[u]){if(v==fath||v==son[u])continue;dfs2(v,u,v);}}
public:void build(int rt){dfs1(rt,0,1);dfs2(rt,0,rt);for(int i=1;i<=n;i++){tr[co[i]].insert(id[i]);}}int find(int u,int v,int c){//查找路径(u,v)上是否有颜色c的节点int res=0;while(top[u]!=top[v]){if(dep[top[u]]<dep[top[v]])swap(u,v);res+=tr[c].query_num(id[top[u]],id[u]);u=fa[top[u]];}if(id[u]>id[v])swap(u,v);res+=tr[c].query_num(id[u],id[v]);return res>=1;}
}tre;
int main(){ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);cin>>n>>m;for(int i=1;i<=n;i++)cin>>co[i];for(int i=1;i<=n-1;i++){int x,y;cin>>x>>y;e[x].push_back(y);e[y].push_back(x);}tre.build(1);//记得建树！！！for(int i=1;i<=m;i++){int x,y,z;cin>>x>>y>>z;cout<<tre.find(x,y,z);}return 0;
}