1. 程式人生 > 其它 >樹上莫隊 洛谷SP10707 COT2 - Count on a tree II

樹上莫隊 洛谷SP10707 COT2 - Count on a tree II

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 typedef long long ll;
  4 const int N=200200;
  5 int dep[N],fa[N][30],color[N],bin[N],blg[N],ord[N],in[N],out[N],ans[N],h[N<<1],cnt[N],vis[N],tot,id,now;
  6 struct event
  7 {
  8     int l,r,pos,lca;
  9 }q[N];
 10 
 11 struct edge
 12
{ 13 int v,next; 14 }e[N<<1]; 15 16 void addedge(int u,int v) 17 { 18 e[++tot]=(edge){v,h[u]}; 19 h[u]=tot; 20 } 21 void build(int n) 22 { 23 int sz=sqrt(n); 24 for(int i=1;i<=n;i++)blg[i]=(i-1)/sz+1; 25 } 26 static bool cmp(event &a,event &b) 27 {
28 return (blg[a.l]^blg[b.l])?(blg[a.l]<blg[b.l]):((blg[a.l]&1)?a.r<b.r:a.r>b.r); 29 } 30 inline ll read() 31 { 32 ll x=0,f=1;char ch=getchar(); 33 while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();} 34 while(ch>='0'&&ch<='9'){x=(x<<1)+(x<<3
)+(ch^48);ch=getchar();} 35 return x*f; 36 } 37 void dfs(int u,int p) 38 { 39 40 in[u]=++id; 41 ord[id]=u; 42 dep[u]=dep[p]+1; 43 fa[u][0]=p; 44 for(int i=1;(1<<i)<dep[u];i++)fa[u][i]=fa[fa[u][i-1]][i-1]; 45 for(int i=h[u];i;i=e[i].next) 46 { 47 int v=e[i].v; 48 if(v==p)continue; 49 dfs(v,u); 50 } 51 out[u]=++id; 52 ord[id]=u; 53 } 54 int lca(int u,int v) 55 { 56 if(dep[u]>dep[v])swap(u,v); 57 for(int i=20;i>=0;i--)if(dep[fa[v][i]]>=dep[u])v=fa[v][i]; 58 if(u==v)return u; 59 for(int i=20;i>=0;i--)if(fa[v][i]!=fa[u][i]){v=fa[v][i];u=fa[u][i];} 60 return fa[u][0]; 61 } 62 void change(int node) 63 { 64 if(vis[node]) 65 { 66 if((--cnt[color[node]])==0)now--; 67 } 68 else 69 { 70 if((++cnt[color[node]])==1)now++; 71 } 72 vis[node]^=1; 73 } 74 int main() 75 { 76 int n=read(),m=read(); 77 for(int i=1;i<=n;i++)color[i]=bin[i]=read(); 78 for(int i=1;i<n;i++) 79 { 80 int u=read(),v=read(); 81 addedge(u,v); 82 addedge(v,u); 83 } 84 dfs(1,0); 85 86 for(int i=1;i<=m;i++) 87 { 88 int l=read(),r=read(); 89 q[i].pos=i; 90 int lca_=lca(l,r); 91 if(in[l]>in[r])swap(l,r); 92 if(l==lca_) 93 { 94 q[i].l=in[l]; 95 q[i].r=in[r]; 96 q[i].lca=0; 97 } 98 else 99 { 100 q[i].l=out[l]; 101 q[i].r=in[r]; 102 q[i].lca=lca_; 103 } 104 105 } 106 build(n*2);//分塊要考慮尤拉序,不然會T; 107 sort(q+1,q+1+m,cmp); 108 sort(bin+1,bin+1+n); 109 int bcnt=unique(bin+1,bin+1+n)-bin-1; 110 for(int i=1;i<=n;i++)color[i]=lower_bound(bin+1,bin+1+bcnt,color[i])-bin; 111 int l=1,r=0; 112 for(int i=1;i<=m;i++) 113 { 114 int ql=q[i].l,qr=q[i].r,lca_=q[i].lca; 115 while(l<ql)change(ord[l++]); 116 while(l>ql)change(ord[--l]); 117 while(r<qr)change(ord[++r]); 118 while(r>qr)change(ord[r--]); 119 if(lca_)change(lca_); 120 ans[q[i].pos]=now; 121 if(lca_)change(lca_); 122 } 123 for(int i=1;i<=m;i++)printf("%d\n",ans[i]); 124 return 0; 125 }