[模板] 点分治

树上距离为k的点对是否存在.

///洛谷P3806

给定一颗n个结点的无根树,有m次询问,每次询问树上距离为k的点对是否存在.

思想:用桶记录路径,判断是否存在距离为k的点对

/**

使用方法:

调用DFZ.sovle()函数后答案存入ans[]数组

模块说明:

1.求树的重心函数 Root.getroot()

输入:

使用Root.getroot(u,fa,sum),u为当前结点,fa为u结点的父亲结点,sum是当前连通块的大小

输出:

返回重心结点编号rt

2.计算所有结点到根节点的距离函数 CalDis.caldis()

输入:

使用CalDis.caldis(int u,int fa),u为当前结点,fa为u结点的父亲结点

输出:

得到di[]数组,di[]数组大小为tp,存有所有基本路径的长度,还有dis[]数组,存的是结点u到当前根节点的长度

3.计算合法路径函数 SovleDis.sovle()

输入:

使用SovleDis.sovle(int u),u为当前结点

输出:

得到当前ans[]数组,,存有每次更新的答案信息

4.点分治函数 DFZ.dfz()

输入:

使用DFZ.dfz(int u),u为当前结点

输出:

递归各结点

5.点分治调用 DFZ.solve()

输入:

无

输出:

得到最终ans[]数组,存有答案信息

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 const int amn=1e5+5,inf=1e9;
  4 int n,m,K[amn];
  5 
  6 ///链式前向星存图
  7 int head[amn],etot;
  8 struct edge{
  9     int nxt,v,w;
 10 }eg[amn];
 11 void add(int u,int v,int w){
 12     eg[++etot]={head[u],v,w};
 13     head[u]=etot;
 14 }
 15 
 16 int vis[amn];
 17 
 18 ///求树的重心
 19 class Root{
 20 public:
 21     int siz[amn],maxt[amn],rt;
 22     void calsiz(int u,int fa,int sum){
 23         siz[u]=1;
 24         maxt[u]=0;
 25         for(int i=head[u];i;i=eg[i].nxt){
 26             int v=eg[i].v;
 27             if(vis[v]||v==fa)continue;
 28             calsiz(v,u,sum);
 29             siz[u]+=siz[v];
 30             maxt[u]=max(maxt[u],siz[v]);
 31         }
 32         maxt[u]=max(maxt[u],sum-siz[u]);
 33         if(maxt[u]<maxt[rt])rt=u;
 34     }
 35     void getroot(int u,int fa,int sum){
 36         rt=0;
 37         maxt[rt]=inf;
 38         calsiz(u,fa,sum);
 39     }
 40 };
 41 
 42 ///求基本路径dis
 43 class CalDis{
 44     public:
 45     int dis[amn],di[amn],tp;
 46     void caldis(int u,int fa){
 47         if(dis[u]>(int)1e7)return;
 48         di[++tp]=dis[u];
 49         for(int i=head[u];i;i=eg[i].nxt){
 50             int v=eg[i].v,w=eg[i].w;
 51             if(vis[v]||v==fa)continue;
 52             dis[v]=dis[u]+w;
 53             caldis(v,u);
 54         }
 55     }
 56     void init(int v,int w){
 57         tp=0;
 58         dis[v]=w;
 59     }
 60 };
 61 
 62 ///判断路径
 63 bool jg[(int)1e7+1];
 64 int ans[amn];
 65 class SovleDis{
 66     public:
 67     CalDis cd;
 68     queue<int> bk;
 69     void sovle(int u){
 70         jg[0]=1;
 71         bk.push(0);
 72         for(int i=head[u];i;i=eg[i].nxt){
 73             int v=eg[i].v,w=eg[i].w;
 74             if(vis[v])continue;
 75             cd.init(v,w);
 76             cd.caldis(v,u);
 77             for(int j=1;j<=cd.tp;j++){
 78                 for(int k=1;k<=m;k++){
 79                     if(K[k]>=cd.di[j])ans[k]+=jg[K[k]-cd.di[j]];
 80                 }
 81             }
 82             for(int j=1;j<=cd.tp;j++){
 83                 jg[cd.di[j]]=1;
 84                 bk.push(cd.di[j]);
 85             }
 86         }
 87         while(bk.size()){
 88             jg[bk.front()]=0;
 89             bk.pop();
 90         }
 91     }
 92 };
 93 
 94 ///点分治
 95 class DFZ{
 96 public:
 97     Root rt;
 98     SovleDis s;
 99     void dfz(int u){
100         vis[u]=1;
101         s.sovle(u);
102         for(int i=head[u];i;i=eg[i].nxt){
103             int v=eg[i].v;
104             if(vis[v])continue;
105             rt.getroot(v,u,rt.siz[v]);
106             dfz(rt.rt);
107         }
108     }
109     void sovle(){
110         rt.getroot(1,-1,n);
111         dfz(rt.rt);
112     }
113 };
114 
115 
116 int main(){
117     DFZ df;
118     int a,b,c;
119     scanf("%d%d",&n,&m);
120     for(int i=1;i<=n-1;i++){
121         scanf("%d%d%d",&a,&b,&c);
122         add(a,b,c);
123         add(b,a,c);
124     }
125     for(int i=1;i<=m;i++){
126         scanf("%d",&K[i]);
127     }
128     df.sovle();
129     for(int i=1;i<=m;i++){
130         if(ans[i])printf("AYE
");
131         else printf("NAY
");
132     }
133 }

树上距离小于等于K的点对数量

///洛谷P4178

给定一棵 n 个节点的树，每条边有边权，求出树上两点距离小于等于 k 的点对数量。

思想:用容斥和双指针记录路径,计算距离小于等于k的点对有多少个.

/**

使用方法:

调用DFZ.sovle()函数后答案存入ans

模块说明:

1.求树的重心函数 Root.getroot()

输入:

使用Root.getroot(u,fa,sum),u为当前结点,fa为u结点的父亲结点,sum是当前连通块的大小

输出:

返回重心结点编号rt

2.计算所有结点到根节点的距离函数 CalDis.caldis()

输入:

使用CalDis.caldis(int u,int fa),u为当前结点,fa为u结点的父亲结点

输出:

得到di[]数组,di[]数组大小为tp,存有所有基本路径的长度,还有dis[]数组,存的是结点u到当前根节点的长度

3.计算合法路径函数 SovleDis.sovle()

输入:

使用SovleDis.sovle(int u,int fa,int w),u为当前结点,fa为u结点的父亲结点,w为fa到u的边权

输出:

得到当前ans,存有每次更新的答案信息

4.点分治函数 DFZ.dfz()

输入:

使用DFZ.dfz(int u),u为当前结点

输出:

递归各结点

5.点分治调用 DFZ.solve()

输入:

无

输出:

得到最终ans,存有答案信息

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 const int amn=1e5+5,inf=1e9;
  4 int n,m,K;
  5 
  6 ///链式前向星存图
  7 int head[amn],etot;
  8 struct edge{
  9     int nxt,v,w;
 10 }eg[amn];
 11 void add(int u,int v,int w){
 12     eg[++etot]={head[u],v,w};
 13     head[u]=etot;
 14 }
 15 
 16 int vis[amn];
 17 
 18 ///求树的重心
 19 class Root{
 20 public:
 21     int siz[amn],maxt[amn],rt;
 22     void calsiz(int u,int fa,int sum){
 23         siz[u]=1;
 24         maxt[u]=0;
 25         for(int i=head[u];i;i=eg[i].nxt){
 26             int v=eg[i].v;
 27             if(vis[v]||v==fa)continue;
 28             calsiz(v,u,sum);
 29             siz[u]+=siz[v];
 30             maxt[u]=max(maxt[u],siz[v]);
 31         }
 32         maxt[u]=max(maxt[u],sum-siz[u]);
 33         if(maxt[u]<maxt[rt])rt=u;
 34     }
 35     void getroot(int u,int fa,int sum){
 36         rt=0;
 37         maxt[rt]=inf;
 38         calsiz(u,fa,sum);
 39     }
 40 };
 41 
 42 ///求基本路径dis
 43 class CalDis{
 44     public:
 45     int dis[amn],di[amn],tp;
 46     void caldis(int u,int fa){
 47         if(dis[u]>K)return;
 48         di[++tp]=dis[u];
 49         for(int i=head[u];i;i=eg[i].nxt){
 50             int v=eg[i].v,w=eg[i].w;
 51             if(vis[v]||v==fa)continue;
 52             dis[v]=dis[u]+w;
 53             caldis(v,u);
 54         }
 55     }
 56     void init(int v,int w){
 57         tp=0;
 58         dis[v]=w;
 59     }
 60 };
 61 
 62 ///判断路径
 63 int ans;
 64 class SovleDis{
 65     public:
 66     CalDis cd;
 67     int sovle(int u,int fa,int w){
 68         cd.init(u,w);
 69         cd.caldis(u,fa);
 70         sort(cd.di+1,cd.di+1+cd.tp);
 71         int l=1,r=cd.tp,ans=0;
 72         while(l<r){
 73             if(cd.di[l]+cd.di[r]<=K){
 74                 ans+=r-l;
 75                 l++;
 76             }
 77             else r--;
 78         }
 79         return ans;
 80     }
 81 };
 82 
 83 ///点分治
 84 class DFZ{
 85 public:
 86     Root rt;
 87     SovleDis s;
 88     void dfz(int u){
 89         vis[u]=1;
 90         ans+=s.sovle(u,-1,0);
 91         for(int i=head[u];i;i=eg[i].nxt){
 92             int v=eg[i].v,w=eg[i].w;
 93             if(vis[v])continue;
 94             ans-=s.sovle(v,u,w);
 95             rt.getroot(v,u,rt.siz[v]);
 96             dfz(rt.rt);
 97         }
 98     }
 99     void sovle(){
100         rt.getroot(1,-1,n);
101         dfz(rt.rt);
102     }
103 };
104 
105 
106 int main(){
107     DFZ df;
108     int a,b,c;
109     scanf("%d",&n);
110     for(int i=1;i<=n-1;i++){
111         scanf("%d%d%d",&a,&b,&c);
112         add(a,b,c);
113         add(b,a,c);
114     }
115     ans=0;
116     scanf("%d",&K);
117     df.sovle();
118     printf("%d
",ans);
119 }

时间/空间复杂度

一．点分治

	总体	caldis()	求重心
时间复杂度	O(nlogn)	O(n)	O(nlogn)
空间复杂度	O(nlogn)	O(n)	O(nlogn)
备注		求基本路径

习题

一．点分治