Skip to main content

Solving CodeChef June Long Challenge Problem "PRMQ" using Segment Tree

Hello folks. Hope you are having a good time. I wanted to share my solution to CodeChef June '17 problem  PRMQ. Though there might be many solutions possible to this problem, I will walk you through a Segment-Tree solution for this.

Approach

There are nearly 78,000 prime numbers between 1 and 10^6, so in worst case a number in this query range L to R can have maximum of 8 prime numbers as the number of factors. Why 8 see 2*3*5*7*11*13*17*19 > 10^6 .

And the root node can have maximum of 78,000 nodes with each prime number in the leaf node merging to form the root node.

Now let us think what we should store in a node of segment tree. We may need the following information:-
1. The element of the array so as to compute their prime factors.

2. Some way to combine the factors of each element as key and their count as value,
suppose we are having 2 elements say 8 and 9 so when we merge these node in the tree, we have got 2 as a factor with count 3 and 3 as a factor with count 2 in the merged node.

We now have count of respective prime factors stored in sorted manner in the merged node. If we iterate through every valid value in range X to Y,  we will be iterating over at min(Y-X+1 , number of prime numbers as factors) * log (N), where N is number of elements in array.
This is pretty costly and motivates us to store prefix count of prime factors as well, isn't it?

3. So, now we also have to store prefix sum of prime factors in each node.
So , all  in all we have the node structure as following ,

Node Structure

1
2
3
4
5
struct node{
 int val;
 vector<pi>M;
 vector<pi>Prefix;
};
the val is the array element , M is like a map but implemented as pair of <int,int> as we have prime factor as key and their count as value, Prefix is same as M but it stores prefix sum of factors instead.

Sieve of Eratosthenes

So now lets see how we are going to build the tree and merge the nodes.
For that first we need to implement Sieve of Eratosthenes to precompute prime factors in the range 1 to 10^6.
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
//sieve of Eratosthenes

void Sieve(){
 int N = 1e6+1;
 bool isPrime[N];
 fill(isPrime,isPrime+N,true);
 for(int i=2;i*i<=N;i++){
  if(isPrime[i]){
   for(int j=i*i;j<=N;j+=i){
    isPrime[j] = false;
   }
  }
 }
 for(int i=2;i<=N;i++){
  if(isPrime[i])
   Primes.pb(i);
 }
}

Segment Tree Construction

Now comes the crucial part of the building the tree. Let's see how we store the leaf nodes in the segment tree. I generally use the iterative implementation from Al.Cash's legendary blog Easy and Efficient Segment Tree.

We declare the Segment Tree globally as following:-
1
2
int SZ = 1e5+1;
vector<node>t(2*SZ);

Filling the leaf nodes

Now we fill the leaf nodes of the trees , We use the Primes vector from the previous sieve() function. and store the required things in the following manner. This part is self explanatory.
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
int main(){
 Sieve();
 int N;
 scanf("%d",&N);
 for(int i=0;i<N;i++){
  int element;
  ini(element);
  int copy = element;
  int prefixsum = 0 ;
  t[N+i].val = element;
  for(int j=0;;j++){
   if(Primes[j] > sqrt(copy)){
    break;
   }
   int cnt=0;
   
   while(element >=1 and element%Primes[j] == 0){
    element/=Primes[j];
    cnt++;
   }
   if(cnt!=0){
    prefixsum += cnt;
    t[N+i].M.pb({Primes[j],cnt});
    t[N+i].Prefix.pb({Primes[j],prefixsum});
   }
  }
  if(element>=2){
   t[N+i].M.pb({element,1});
   t[N+i].Prefix.pb({element,prefixsum+1});
  }
 }
 build(N);
}

Building the tree and merging the nodes 

Now, let us see how we build the tree and merge the nodes. For merging the node, I use the idea how we merge in Merge-Sort, so merging the nodes can be done in linear times the number of nodes in each leaf for level just above leaf and at root maximum 78,000 nodes need to be merged.
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
void build(int N){
 for(int i=N-1; i> 0 ;i--){
  
  node left = t[i<<1];
  node right = t[i<<1|1];
  
  t[i].val = left.val + right.val;
  int n = left.M.size();
  int m = right.M.size();
  
  int p=0,q=0;
  int prefixsum  = 0 ;
  while(p<n and q<m){
   if(left.M[p].ff < right.M[q].ff){
    t[i].M.pb(left.M[p]);
    prefixsum += left.M[p].ss;
    t[i].Prefix.pb({left.M[p].ff , prefixsum});
    p++;
   }
   else if(left.M[p].ff > right.M[q].ff){
    t[i].M.pb(right.M[q]);
    prefixsum += right.M[q].ss;
    t[i].Prefix.pb({right.M[q].ff , prefixsum});
    q++;
   }
   else{
    t[i].M.pb({left.M[p].ff,left.M[p].ss+right.M[q].ss});
    prefixsum += left.M[p].ss + right.M[q].ss;
    t[i].Prefix.pb({left.M[p].ff , prefixsum});
    p++;
    q++;
   }
  }
  while(p<n){
   prefixsum += left.M[p].ss;
   t[i].M.pb({left.M[p]});
   t[i].Prefix.pb({left.M[p].ff , prefixsum});
   p++;
  }
  while(q<m){
   prefixsum += right.M[q].ss;
   t[i].M.pb({right.M[q]});
   t[i].Prefix.pb({right.M[q].ff , prefixsum});
   q++;
  }
  
 }
}
Above code is complex but algorithm of merging is easy , you can follow with an example using any 4 node at the leaves.

Querying the Segment Tree

Querying the tree is straight forward, we keep on merging the nodes which are required (merging is clearly explained in the Al.Cash's blog, please refer that). It will explain you why we merge the way we merge.
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
int query(int N, int L, int R , int X, int Y){
 
 int ANS  = 0;
 for(L += N , R += N ;  L < R ; L >>= 1 , R >>= 1){
  if(L&1){
   int leftI =  Greater_equal(t[L].Prefix , X);
   int rightI =  Lesser_equal(t[L].Prefix , Y);
   if(leftI != -1 and rightI != -1){
    if(leftI-1>=0){
     ANS += (t[L].Prefix[rightI].ss - t[L].Prefix[leftI-1].ss);    
    }
    else{
     ANS += (t[L].Prefix[rightI].ss);
    }
   }
   L++;
  }
  
  if(R&1){
   --R;
   int leftI =  Greater_equal(t[R].Prefix , X);
   int rightI =  Lesser_equal(t[R].Prefix , Y);
   if(leftI != -1 and rightI != -1){
    if(leftI-1>=0){
     ANS += (t[R].Prefix[rightI].ss - t[R].Prefix[leftI-1].ss);    
    }
    else{
     ANS += (t[R].Prefix[rightI].ss);
    }
   }
   
  }
 }
 return ANS;
 
}

Custom Utility Functions

Custom utility functions are required in the Query function, because we need to find the prefix sum, for that we need to find the prime factor in the node greater than or equal to X and the prime factor less than or equal to the Y. So I wrote the functions by myself,  you can use library functions too, but you need to write custom comparator anyway, so I figured it's better to write on my own. 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
int Greater_equal(vector<pi> &a, int B){
 int lo = 0;
 int hi = a.size()-1;
 
 int index = -1;
 while(lo<=hi){
  int mid = lo + (hi-lo)/2;
  
  if(a[mid].ff < B){
   lo = mid + 1;
  }
  else{
   index = mid;
   hi = mid -1;
  }
 }
 return index;
}
 
int Lesser_equal(vector<pi> &a, int B){
 int lo = 0;
 int hi = a.size()-1;
 
 int index = -1;
 while(lo<=hi){
  int mid = lo + (hi-lo)/2;
  
  if(a[mid].ff > B){
   hi  = mid - 1;
  }
  else{
   index = mid;
   lo = mid + 1;
  }
 }
 return index;
}
Now we combine all the bits we have seen so far in the final solution. If you feel you understand the approach but don't understand some parts,  I recommend you to trace complete code with some simple examples first for segment tree then for this code.

Here is the link to the AC solution check it out Solution.
Please share this blog and comment your feed backs.

Comments

  1. Nice explanation!!

    ReplyDelete
  2. nice work.. :) and that segment tree blog is awesome.

    ReplyDelete
    Replies
    1. thank you , truly it is !

      Delete
    2. there is no use of val so we can construct struct without val :)

      Delete
    3. yeah sure, but it doesn't hurt, just in case you need it. :)

      Delete
  3. Sir can we use stl maps instead of vector of pairs. merging of stl maps can be done easily.

    ReplyDelete
    Replies
    1. If you check my previous submissions, I earlier tried maps and somehow in building trees it takes more than 1 sec to build the tree, due to log(N) complexity for insertion and other overheads, so that gives TLE in one test case in last task. so i guess maps won't work, but you can try to optimize that.

      Delete
  4. What is the memory complexity of the segment tree?

    ReplyDelete
    Replies
    1. Let's see the memory complexity, at root node we can have maximum 78000 nodes, and below that we will be having 2 * 39,000 , below that 4*19500...and son on till height of tree.
      so N + 2*N/2 + 4*N/4 + .........+ log(78000)*78000/log(78,000) = 78,000*log(78000).

      This memory is just for the one vector and we have two of them. Otherwise if we have duplicate copies of factors in leaf node, we can have maximum 8 of them in a node at leaves even in that case we will have same N*log(N) space complexity.

      Delete
  5. Nice explanation!! but i want to know how u build such a theme for blogspot. it's really nice

    ReplyDelete
    Replies
    1. I have used one of the standard themes for blogspot.

      Delete
  6. Wow, you're using the compact segment tree from leaves. I like that!

    ReplyDelete
  7. Hey, I followed the exact same approach but got TLE. All complexities are same.
    https://www.codechef.com/viewsolution/14389097
    Please check it for me and tell whats wrong. Couldn't figure it out

    ReplyDelete
  8. Doing the same but getting TLE in subtask-3, Please check it out this link - https://www.codechef.com/viewsolution/15155710

    ReplyDelete
  9. Thanks for the explanation. The official editorial was really frustrating to understand.

    ReplyDelete
  10. very nice explanation sir!

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

    ReplyDelete
  12. You can make the article awesome by putting snapshots of tree after every function or operation.
    By the way good article.

    ReplyDelete

Post a Comment

Popular posts from this blog

Topcoder SRM 684 DivFreed2 Div 2 level 2

Hello all, today I want to share a problem which comes in many forms of counting the number of sequences subject to some constraints on the elements of the sequence. The problem goes like following. link to problem statement ------------------------------------------------------------------- Problem Statement Hero likes some arrays. The arrays he likes are the arrays that have all of the following properties: The length of the array is n . Each element is an integer between 1 and k , inclusive. Whenever A and B are two consecutive elements of the array (in this order), we have (A <= B) or (A mod B != 0). For example, suppose n =4 and k =7. Hero will like the array {1,7,7,2} because it has the right length, all elements are in the correct range, 1 <= 7, 7 <= 7, and 7 mod 2 != 0. Hero will not like the array {4,4,4,2}. You are given the ints n and k . Let X be the number of different arrays Hero likes. Compute and return the value (X mo

Modular Exponentiation

Just wanted to introduce you to very basic concept of Modular Exponentiation through a problem,  Big Mod . We need to calculate A^B % M for large inputs. For a second let's forget about A^B  % M and think about A^B only and say A^B lies within the range of long long int and B is less than 10^7. A basic brute force approach to calculate A^B will be to loop through 1 to B and multiply A every time. The time complexity of this approach is O(n). This is perfect on current constraints but what happens if I change the constraint on B and raise it to 10^12 or so. Now  this approach will give a TLE. An efficient approach is needed. O(n) is not sufficient. We will use a recursive approach for this. Suppose I want to calculate A^B but if B is even, I can calculate A^(B/2) and multiply it to itself, and the way to calculate A^(B/2) is to calculate A^(B/4) and so on you get the idea. But what if B is not even, no problem,  B-1 will be even do the same except replace B with B-1 and separ