[CodeVS4633][Mz]树链剖分练习

Posted 2020-09-28 skylee的OI博客

思路：

轻重链剖分+线段树。

  1 #include<cstdio>
  2 #include<vector>
  3 #include<cstring>
  4 const int N=100001;
  5 std::vector<int> e[N];
  6 inline void add_edge(const int u,const int v) {
  7     e[u].push_back(v);
  8     e[v].push_back(u);
  9 }
 10 int par[N]={0,1},dep[N]={0},size[N]={0};
 11 void dfs1(const int x) {
 12     size[x]=1;
 13     for(unsigned int i=0;i<e[x].size();i++) {
 14         if(e[x][i]==par[x]) continue;
 15         par[e[x][i]]=x,dep[e[x][i]]=dep[x]+1;
 16         dfs1(e[x][i]);
 17         size[x]+=size[e[x][i]];
 18     }
 19 }
 20 int top[N]={0},dfn[N],sz=0;
 21 void dfs2(const int x) {
 22     if(!top[x]) top[x]=x;
 23     dfn[x]=++sz;
 24     if(e[x].size()==1&&x!=1) return;
 25     int v=0;
 26     for(unsigned int i=0;i<e[x].size();i++) {
 27         if(e[x][i]==par[x]) continue;
 28         if(size[e[x][i]]>size[v]) v=e[x][i];
 29     }
 30     top[v]=top[x];
 31     dfs2(v);
 32     for(unsigned int i=0;i<e[x].size();i++) {
 33         if(e[x][i]==par[x]||e[x][i]==v) continue;
 34         dfs2(e[x][i]);
 35     }
 36 }
 37 class SegmentTree {
 38     #define _left <<1
 39     #define _right <<1|1
 40     private:
 41         int val[N<<2],tag[N<<2];
 42         int len(const int l,const int r) {
 43             return r-l+1;
 44         }
 45         void push_down(const int p,const int b,const int e) {
 46             if(!tag[p]) return;
 47             int mid=(b+e)>>1;
 48             tag[p _left]+=tag[p];
 49             tag[p _right]+=tag[p];
 50             val[p _left]+=tag[p]*len(b,mid);
 51             val[p _right]+=tag[p]*len(mid+1,e);
 52             tag[p]=0;
 53         }
 54         void push_up(const int p) {
 55             val[p]=val[p _left]+val[p _right];
 56         }
 57     public:
 58         SegmentTree() {
 59             memset(val,0,sizeof val);
 60             memset(tag,0,sizeof tag);
 61         }
 62         void modify(const int p,const int b,const int e,const int l,const int r) {
 63             if((b==l)&&(e==r)) {
 64                 val[p]+=len(b,e);
 65                 tag[p]++;
 66                 return;
 67             }
 68             push_down(p,b,e);
 69             int mid=(b+e)>>1;
 70             if(l<=mid) modify(p _left,b,mid,l,std::min(mid,r));
 71             if(r>mid) modify(p _right,mid+1,e,std::max(mid+1,l),r);
 72             push_up(p);
 73         }
 74         int query(const int p,const int b,const int e,const int l,const int r) {
 75             if((b==l)&&(e==r)) return val[p];
 76             push_down(p,b,e);
 77             int mid=(b+e)>>1,ans=0;
 78             if(l<=mid) ans+=query(p _left,b,mid,l,std::min(mid,r));
 79             if(r>mid) ans+=query(p _right,mid+1,e,std::max(mid+1,l),r);
 80             return ans;
 81         }
 82 };
 83 SegmentTree t;
 84 inline void swap(int &x,int &y) {
 85     int t;
 86     t=x;
 87     x=y;
 88     y=t;
 89 }
 90 int n;
 91 void modify(int x,int y) {
 92     for(;top[x]!=top[y];x=par[top[x]]) {
 93         if(dep[top[x]]<dep[top[y]]) swap(x,y);
 94         t.modify(1,1,n,dfn[top[x]],dfn[x]);
 95     }
 96     if(dep[x]<dep[y]) swap(x,y);
 97     t.modify(1,1,n,dfn[y],dfn[x]);
 98 }
 99 int query(int x,int y) {
100     int ans=0;
101     for(;top[x]!=top[y];x=par[top[x]]) {
102         if(dep[top[x]]<dep[top[y]]) swap(x,y);
103         ans+=t.query(1,1,n,dfn[top[x]],dfn[x]);
104     }
105     if(dep[x]<dep[y]) swap(x,y);
106     ans+=t.query(1,1,n,dfn[y],dfn[x]);
107     return ans;
108 }
109 int main() {
110     scanf("%d",&n);
111     for(int i=1;i<n;i++) {
112         int x,y;
113         scanf("%d%d",&x,&y);
114         add_edge(x,y);
115     }
116     dfs1(1);
117     dfs2(1);
118     int q;
119     scanf("%d",&q);
120     while(q--) {
121         int op,x,y;
122         scanf("%d%d%d",&op,&x,&y);
123         if(op==1) modify(x,y);
124         if(op==2) printf("%d\n",query(x,y));
125     }
126     return 0;
127 }