913. Query on a tree II
http://www.spoj.com/problems/QTREE2/
コスト付き木が与えられる。
頂点a,b間のコストを答えるクエリと
頂点a,b間のパスのk番目の頂点を答えるクエリを処理しろというような問題。
Heavy-Light Decompositionを用いて、LCAなどを処理するようにしました。
速度はそこまで早くないようです。
もう少しコンパクトな実装にしたいです。
struct Heavylight{ vector<vector<PI> > G; vector<int> sizevis; vector<vector<int> > segments; vector<PI> pos; vector<int> parent; vector<int> segdist; void sizedfs(int v){ queue<int> q; q.push(v); vector<int> near; sizevis = vector<int>(G.size()); while(!q.empty()){ int cv = q.front(); q.pop(); if(sizevis[cv]) continue; sizevis[cv] = 1; near.pb(cv); FOR(it,G[cv]){ int to=it->F; if(sizevis[to]) continue; q.push(to); } } sizevis = vector<int>(G.size()); FORR(it,near){ int ret=1; FOR(ch,G[*it]){ int to=ch->F; if(sizevis[to]) ret += sizevis[to]; else parent[*it] = to; } sizevis[*it] = ret; } } void makesegment(int v){ vector<int> childsegment; segments.push_back(vector<int>(1,v)); vector<int>& currentsegment = segments.back(); while(true){ vector<PI> size; FOR(it,G[v]) if(sizevis[it->F]<sizevis[v]) size.pb(mp(-sizevis[it->F],it-G[v].begin())); sort(ALL(size)); for(int i=1;i<SZ(size);++i){ const PI& e=G[v][size[i].S]; childsegment.pb(e.F); segdist[e.F] = e.S; } if(SZ(size)){ const PI& e = G[v][size[0].S]; int chv = e.F; currentsegment.pb(chv); segdist[chv] = segdist[v] + e.S; v = chv; }else break; } for(int i=0;i<SZ(currentsegment);++i) pos[currentsegment[i]] = mp(segments.size()-1,i); FOR(it,childsegment) makesegment(*it); } void init(int root){ parent = vector<int>(G.size(),-1); pos = vector<PI>(G.size()); segdist = vector<int>(G.size()); sizedfs(root); makesegment(root); } Heavylight(const vector<vector<PI> >& aG,int root=0):G(aG){ init(root); } int lca(int u,int v){ while(true){ if(pos[u].F == pos[v].F) return pos[u].S>pos[v].S?v:u; if(sizevis[segments[pos[u].F][0]] > sizevis[segments[pos[v].F][0]]) swap(u,v); u = parent[segments[pos[u].F][0]]; } } int dist(int u,int v){ int ret=0; while(true){ if(pos[u].F == pos[v].F) return ret + abs(segdist[u]-segdist[v]); if(sizevis[segments[pos[u].F][0]] > sizevis[segments[pos[v].F][0]]) swap(u,v); ret += segdist[u]; u = parent[segments[pos[u].F][0]]; } } int kth(int u,int v,int k){// u=v0,v=vn,return vk int lcav = lca(u,v); while(pos[u].F != pos[lcav].F){ if(pos[u].S>=k) return segments[pos[u].F][pos[u].S-k]; k -= pos[u].S+1; u = parent[segments[pos[u].F][0]]; } if(pos[u].S-pos[lcav].S>=k) return segments[pos[u].F][pos[u].S-k]; k -= pos[u].S - pos[lcav].S; int tv = v; int di = 0; while(pos[tv].F != pos[lcav].F){ di += pos[tv].S+1; tv = parent[segments[pos[tv].F][0]]; } di += pos[tv].S - pos[lcav].S; k = di - k; while(pos[v].F != pos[lcav].F){ if(pos[v].S>=k) return segments[pos[v].F][pos[v].S-k]; k -= pos[v].S+1; v = parent[segments[pos[v].F][0]]; } if(pos[v].S-pos[lcav].S>=k) return segments[pos[v].F][pos[v].S-k]; } }; inline int in(){ int n; scanf("%d",&n); return n; } void solve(){ int n=in(); vector<vector<PI> > G(n); rep(i,n-1){ int a,b,c; a=in()-1; b=in()-1; c=in(); G[a].pb(mp(b,c)); G[b].pb(mp(a,c)); } Heavylight hl(G); char q[10]; while(true){ int a,b; scanf(" %s",q); if(q[1]=='I'){ a=in(); b=in(); printf("%d\n",hl.dist(a-1,b-1)); }else if(q[1]=='T'){ int k; a=in(); b=in(); k=in(); printf("%d\n",hl.kth(a-1,b-1,k-1) + 1); }else break; } } int main(int argc, char *argv[]) { int t=in(); rep(i,t) solve(); return 0; }