MO's Algorithm to find number of elements present in both array

Question

I have 2 arrays, before[N+1](1 indexed) and after[] (subarray of before[]). Now for M Queries, I need to find how many elements of after[] are present in before[] for the given range l,r.

For example:
N = 5
Before: (2, 1, 3, 4, 5)
After: (1, 3, 4, 5)
M = 2
L = 1, R = 5 → 4 elements (1, 3, 4, 5) of after[] are present in between before[1] and before[5]
L = 2, R = 4 → 3 elements (1, 3, 4) of after[] are present in between before[2] and before[4]

I am trying to use MO's algorithm to find this.Following is my code :

using namespace std;

int N, Q;

// Variables, that hold current "state" of computation
long long current_answer;
long long cnt[100500];

// Array to store answers (because the order we achieve them is messed up)
long long answers[100500];
int BLOCK_SIZE;


// We will represent each query as three numbers: L, R, idx. Idx is 
// the position (in original order) of this query.
pair< pair<int, int>, int> queries[100500];


// Essential part of Mo's algorithm: comparator, which we will
// use with std::sort. It is a function, which must return True
// if query x must come earlier than query y, and False otherwise.
inline bool mo_cmp(const pair< pair<int, int>, int> &x,
        const pair< pair<int, int>, int> &y)
{
    int block_x = x.first.first / BLOCK_SIZE;
    int block_y = y.first.first / BLOCK_SIZE;
    if(block_x != block_y)
        return block_x < block_y;
    return x.first.second < y.first.second;
}

// When adding a number, we first nullify it's effect on current
// answer, then update cnt array, then account for it's effect again.
inline void add(int x)
{
    current_answer -= cnt[x] * cnt[x] * x;
    cnt[x]++;
    current_answer += cnt[x] * cnt[x] * x;
}

// Removing is much like adding.
inline void remove(int x)
{
    current_answer -= cnt[x] * cnt[x] * x;
    cnt[x]--;
    current_answer += cnt[x] * cnt[x] * x;
}

int main()
{
    cin.sync_with_stdio(false);
    cin >> N >> Q;    // Q- number of queries
    BLOCK_SIZE = static_cast<int>(sqrt(N));


   long long int before[N+1];   // 1 indexed
   long long int after[]        // subarray

    // Read input queries, which are 0-indexed. Store each query's 
    // original position. We will use it when printing answer.
    for(long long int i = 0; i < Q; i++) {
        cin >> queries[i].first.first >> queries[i].first.second;
        queries[i].second = i;
    }

    // Sort queries using Mo's special comparator we defined.
    sort(queries, queries + Q, mo_cmp);

    // Set up current segment [mo_left, mo_right].
    int mo_left = 0, mo_right = -1;

    for(long long int i = 0; i < Q; i++) {
        // [left, right] is what query we must answer now.
        int left = queries[i].first.first;
        int right = queries[i].first.second;

        // Usual part of applying Mo's algorithm: moving mo_left
        // and mo_right.
        while(mo_right < right) {
            mo_right++;
            add(after[mo_right]);
        }
        while(mo_right > right) {
            remove(after[mo_right]);
            mo_right--;
        }

        while(mo_left < left) {
            remove(after[mo_left]);
            mo_left++;
        }
        while(mo_left > left) {
            mo_left--;
            add(after[mo_left]);
        }

        // Store the answer into required position.
        answers[queries[i].second] = current_answer;
    }

    // We output answers *after* we process all queries.
    for(long long int i = 0; i < Q; i++)
        cout << answers[i] << "\n";

Now the problem is I can't figure out how to define add function and remove function.

Can someone help me out with these functions ?

Answer 1

Note: I'll denote the given arrays as a and b.

1. Let's learn how to add a new position (move right by one). If a[r] is already there, you can just ignore it. Otherwise, we need to add a[r] and add the number of occurrences of b[r] in a so far to the answer. Finally, if b[r] is already in a, we need to add one to the answer. Note that we need two count to arrays to do that: one for the first array and one for the second.

2. We know how to add one position in O(1), so we're almost there. How do we handle deletions?

3. Let's assume that we want to remove a subsegment. We can easily modify the count arrays. But how do we restore the answer? Well, we don't. Your solution goes like this:

   • save the current answer
   • add a subsegment
   • answer the query
   • remove it (we take care about the count arrays and ignore the answer)
   • restore the saved answer

That's it. It would require rebuilding the structure when we move the left pointer to the next block, but it still requires O(N sqrt(N)) time in the worst case.

Note: it might be possible to recompute the answer directly using count arrays when we remove one position, but the way I showed above looks easier too me.