Square Root Decomposition with range updates

Hello guys! Hope you all are doing well.

In Codechef December challenge I came across a very interesting problem which can be solved using square root decomposition. So i am gonna explain the algorithm using that particular problem. Link to problem.

Pre-requisites: properties of XOR operation.
Problem:
There is an array and we have to solve 2 types of queries on it:

• type 1: Given two numbers i and x, the value at index i should be updated to x.
• type 2: Given two numbers i and k, your program should output the total number of magical subarrays with the last index ≤ i in which the xor of all elements is equal to k.

Solution for 50 marks: 
For type 1 query just update the given array in O(1) time.
For type 2 query we can use dp. Make prefix_xor array for the given array. Now we can check where prefix_xor[i] = k in linear time.

// For type-2 query
// a[] is the given array
int ans = 0;
prefix_xor[0] = a[0]
if(a[0] == k){
    ans++;
}
for(int i = 1 ; i < n ; i++){
    prefix_xor[i] = prefix_xor[i-1] ^ a[i];
    if(prefix_xor[i] == k){
        ans++;
    }
}

Linear time for each query is not sufficient to solve this problem, we need to reduce the complexity further. Let's see how can we do that.

Solution for 100 marks:
First of all we will make a prefix_xor array similar to the one made above but only once. After this we will divide the prefix array into buckets each of size sqrt(n). Each bucket is sorted so that for query of type2 we can use binary search. We are also saving our original array(which is given) in b[].

typedef long long ll;
const int N = 100010;

// inside main()
ll b[n], a[n];
for(int i = 0 ; i < n; i++){
    cin >> a[i];
    b[i] = a[i];
    if(i){
        a[i] = (a[i] ^ a[i-1]); // a[] is prefix_xor array
    }
}

vector<ll> buckets[N];
int num = 0;  // num is number of buckets

for(int i = 0 ; i < n; i += (ll)sqrt(n), num++){
    for(int j = i ; j < i + (ll)sqrt(n) && j < n; j++){
        buckets[num].push_back(a[j]);
    }
    sort(all(buckets[num]));
}   

For type-1 query (i, x) in which we need to change only one element and we know that we are storing prefix xor in our buckets, so all the elements after this element needs to be updated also. Here we will use a lazy array which will update all buckets in sqrt(n) time. For each query of type 1 we need to update only one bucket (lets say B) in which given index i lies. Buckets which are after B are the ones in which all elements needs to be updated so, we just update the lazy[] with (lazy[bucket_idx] ^ original_value_at_i ^ new_value_for_i).
Time Complexity: O(sqrt(n)) for each query of type-1. 

int i, new_val;
cin >> i >> new_val;

if(type == 1){
    ll lowest_bucket_index;
    // consider 0-based indexing
    lowest_bucket_index = (i-1) / (ll)sqrt(n);
    ll buffer = (b[i-1] ^ new_val);
    b[idx-1] = y; // update given index with new value
    // Empty the bucket 'B' which contain element at index i
    buckets[lowest_bucket_index].clear();
    ll lastidx = (lowest_bucket_index + 1) * (ll)sqrt(n);
    //------- Update the first bucket 'B'--------
        for(int i = idx - 1; i < lastidx  ; i ++){
            a[i] = (a[i] ^ buffer);
        }
        for(int i = lowest_bucket_index * (ll)sqrt(n); i < lastidx ; i ++){
            buckets[lo_buck_num].pb(a[i]);
        }
        sort(all(buckets[lo_buck_num]));
    //--------------------------------------------
    
    //****** Update remaining buckets ************
    for(int i = lo_buck_num+1 ; i < num ; i ++){
        lazy[i] = (lazy[i] ^ buffer);
    }    
    //********************************************
}

For type-2 Query(idx, y) we need to find the number of indices < idx where prefix xor equals y. For this query we also search one bucket completely(upto index idx) in which idx lies. Buckets before this bucket, contain prefix xor in sorted order, so in each bucket we can do binary search 2 times to look for upper index for y and lower index for y. Finally add (upper_idx - lower_idx + 1) in the final answer for each bucket.
Note: Buckets can be updated only when XORred with lazy[]. So is idea of binary search on the buckets is wrong? No it's not because we will update our search element y instead of updating each element of bucket by XORring it with lazy[bucket_idx].

ll idx, y;
cin >> idx >> y;
ll hi = idx-1;
hi_bucket_idx = hi / (ll)sqrt(n);
ll startidx = hi_bucket_idx * (ll)sqrt(n);
ll ans = 0;
for(int i = startidx; i <= hi ; i ++){
    if((a[i]^lazy[hi_bucket_idx]) == y){
        ans++;
    }
}

// Binary search, for each bucket < hi_bucket_idx
for(int i = 0 ; i < hi_bucket_idx ; i ++){
    ll foo; // upper index
    ll bar; // lower index
    foo = bar = -1;
    ll left, right;
    left = 0;
    ll x = (y ^ lazy[i]);   // update search element y
    right = buckets[i].size()-1;
    while(left <= right){
        ll mid = (left + right) >> 1;
        if((mid == 0 || (x > buckets[i][mid-1])) && buckets[i][mid] == x){
            bar = mid;break;
        }
        else if(x > buckets[i][mid]){
          left = mid+1;
        }
        else{
          right = mid-1;
        }
    }
    left = 0;
    right = buckets[i].size()-1;
    while(left <= right){
        ll mid = (left + right) >> 1;
        if((mid == sz(buckets[i])-1 || (x < buckets[i][mid+1])) && buckets[i][mid] == x){
            foo = mid;break;
        }
        else if(x < buckets[i][mid])
          right = mid-1;
        else
          left = mid+1;
    }
    if(foo >= bar and foo != -1 and bar != -1){
        ans += 1 + foo - bar;
    }
}
cout << ans << endl;

Time Complexity: O(sqrt(n) * log(n)) for each query of type-2.

Thank you guys for reading this post. If you find anything wrong in this post please comment below and i hope you learnt something usefull from this post.
- Happy Coding