Friday, December 10, 2010

SPOJ QTREE2

I'm back after a long break; longer than I anticipated. I was reading a Topcoder tutorial on Range Minimum Query (RMQ) and Lowest Common Ancestor (LCA) problems. Here's a comparison table of the algorithms presented in the tutorial -



RMQ Algorithms

AlgorithmTime Complexity
(pre-processing, per-query)
Memory ComplexityCompatibility with dynamic data
SQRT(N) Segment AlgorithmO(N), O(sqrt(N))O(N)Compatible
Sparse TableO(NlogN), O(1)O(NlogN)Not compatible
Segment TreeO(N), O(logN)O(N)Compatible

LCA Algorithms

AlgorithmTime Complexity
(pre-processing, per-query)
Memory Complexity
SQRT(H) Segment AlgorithmO(N), O(sqrt(N))O(N)
Dynamic programming
(similar to Sparse Table)
O(NlogN), O(logN)O(NlogN)

LCA problems can be converted to restricted RMQ problems where consecutive elements differ by at most 1 in O(N) time. RMQ problems can also be converted to LCA problem using Cartesian tree in O(N) time. So, in turn, we can convert any RMQ problem to a restricted RMQ problem by first converting it to LCA and then to RMQ again. Restricted RMQ can be solved using a <O(N), O(1)> algorithm.


This problem's link was given at the end of the tutorial as practice problem.

In this problem
you are given a weighted tree with N nodes (N<=10000) followed by some queries. Queries can be of 2 types - DIST and KTH. In a DIST query, you have to calculate the distance between 2 nodes and in a KTH query, we have to print the k-th node on the path from a particular node to another.
Firstly, I have selected a random node as the root of the tree and then ran a DFS to calculate distance to each node from the root. Now, distance between 2 nodes can be easily calculated if we know their LCA.
Dist(a,b) = Dist(root, a) + Dist(root,b) - 2 * Dist(root, LCA)
 I've used the dynamic programming solution similar to Sparse Table for finding the LCA. In this algorithm, we  store the 2^j-th ancestor of each node i. Read the tutorial for details.

Now, when finding the k-th node on the path from a to b, we need to first determine on which segment this k-th node appears

  • a to LCA
  • LCA to b
If it is on the 2nd segment, we can just swap the nodes and change the value of k accordingly. Then we're left with finding k-th ancestor of a particular node. This can be done utilizing the ancestor array that we pre-computed.

C++ source code below - 

/*
Problem Name : QTREE2
Problem Number : SPOJ 913
Link : https://www.spoj.pl/problems/QTREE2/
Problem Type : Search, LCA
Difficulty : 7 / 10
Interest : 9 / 10
Complexity : O(NlogN), O(logN)
Date : December 10, 2010
*/

#include <stdio.h>
#include <vector>
#include <time.h>
#include <string>
#include <assert.h>
#include <algorithm>
using namespace std;

#define rep(i,n) for(i=0;i<(n);i++)
#define MAXN 10000
#define MAXLOGN 14
//#define i64 __int64
#define i64 long long

vector< pair<int,int> > adj[MAXN+1];

int n;
i64 d[MAXN+1];
char flag[MAXN+1];
int root;
int L[MAXN+1];
int father[MAXN+1];
int A[MAXN+1][MAXLOGN+1];
i64 INF = 10000000;

void dfs(int cur) {
 flag[cur] = true;
 int i;
 rep(i,adj[cur].size()) {
  if(!flag[adj[cur][i].first]) {
   d[adj[cur][i].first] = d[cur] + adj[cur][i].second;
   father[adj[cur][i].first] = cur;
   L[adj[cur][i].first] = L[cur] + 1;
   dfs(adj[cur][i].first);
  }
 }
}

vector<string> parse(const string& s,const string& delim=" ") {
 vector<string>res;
 string t;
 for(int i=0;i!=s.size();i++) {
  if(delim.find(s[i]) != string::npos) {
   if(!t.empty()) {
    res.push_back(t);
    t="";
   }
  }
  else t+=s[i];
 }
 if(!t.empty()) res.push_back(t);
 return res;
}

void pre_process() {
 int i,j;
 for(i=1;i<=n;i++) for(j=0;(1<<j)<=n;j++) A[i][j] = -1;
 for(i=1;i<=n;i++) A[i][0] = father[i];
 for(j=1;(1<<j)<=n;j++) for(i=1;i<=n;i++) if(A[i][j-1] != -1) A[i][j] = A[A[i][j-1]][j-1];
}

int getLCA(int p, int q) {
 int i,log;
 if(L[p] < L[q]) swap(p,q);
 for(log=1;(1<<log)<=L[p];log++);
 log--;
 for(i=log;i>=0;i--) if(L[p] - (1<<i) >= L[q]) p = A[p][i];
 if(p == q) return p;
 for(i=log;i>=0;i--) {
  if(A[p][i] != -1 && A[p][i] != A[q][i]) {
   p = A[p][i];
   q = A[q][i];
  }
 }
 return father[p];
}


int getKth(int p, int q, int k, int LCA) {
 int d,log,i;
 if(LCA == p) {
  d = L[q] - L[p] + 1;
  swap(p,q);
  k = d - k + 1;
 }
 else if(LCA == q) ;
 else {
  if(L[p] - L[LCA] + 1 < k) {
   d = L[p] + L[q] - 2 * L[LCA] + 1;
   k = d - k + 1;
   swap(p,q);
  }
 }
 for(log=1;(1<<log)<=L[p];log++);
 log--;
 k--;
 for(i=log;i>=0;i--) if((1<<i) <= k) { 
  p = A[p][i];
  k -= (1<<i);
 }
 return p;
}

int main() {
 int T;
 int i,a,b,c,k;
 i64 res;
 char s[100];
 int LCA;
 INF *= INF;
 vector<string> vs;
 scanf(" %d",&T);
 srand(time(NULL));
 while(T--) {
  scanf(" %d",&n);
  for(i=1;i<=n;i++) adj[i].clear(), d[i] = INF, flag[i] = false;
  for(i=1;i<n;i++) {
   scanf(" %d %d %d",&a,&b,&c);
   adj[a].push_back(make_pair(b,c));
   adj[b].push_back(make_pair(a,c));
  }
  root = rand() % n + 1;
  d[root] = 0;
  father[root] = -1;
  L[root] = 0;
  dfs(root);

  pre_process();

  gets(s);
  while(gets(s)) {
   if(s[1] == 'O') break;
   vs = parse(s);
   a = atoi(vs[1].c_str());
   b = atoi(vs[2].c_str());
   LCA = getLCA(a,b);
   assert(LCA >= 1 && LCA <= n);
   if(vs[0] == "DIST") {
    if(LCA == a || LCA == b) res = abs(d[a] - d[b]);
    else res = d[a] + d[b] - 2 * d[LCA];
   }
   else if(vs[0] == "KTH") {
    k = atoi(vs[3].c_str());
    res = getKth(a,b,k,LCA);
   }
   else assert(0);
   printf("%lld\n",res);
  }
  printf("\n");
 }
 return 0;
}

14 comments:

  1. Nice one Sanny Vai. I have been solving spoj for last few months and there are lots of cool data structure problems.
    Another question, what is the syntax highlighter u r using?

    ReplyDelete
  2. I'm using this as syntax highlighter -
    http://code.google.com/p/google-code-prettify/

    ReplyDelete
  3. hey i tried to run it......
    it is giving me segmentation fault
    :(

    ReplyDelete
  4. Is the seg fault caused by one of the asserts?

    ReplyDelete
  5. QTREE2 te wa ar tle khete khete mara jachilam. Apnar code dekhe bug gula dur korte pere AC pelam. Thanks vaia :)

    ReplyDelete
  6. My solution: runs in 0.66s. Still some optimizations are possible, but this is quite a simple solution http://codepad.org/FaijG9qR

    ReplyDelete
  7. http://code.google.com/p/google-code-prettify/

    ReplyDelete
  8. This comment has been removed by the author.

    ReplyDelete
  9. okay i am new to this kind of data structure
    can anyone explain if Dist(a,b) = Dist(a, LCA(a,b)) + Dist(b,LAC(a,b))
    will work or not for this problem

    thanks in advance :)

    ReplyDelete
  10. more-over why use Sparse Table for O(sqrt(n)) when segment tree
    can do the job in O(log(n))

    ReplyDelete
  11. Sir,can you please do some commenting :)

    ReplyDelete