How can I solve this using BIT?

Question

I found a nice math problem, but I still can't solve it, I tried to find one solution using google and found that it can be solve using the Binary Indexed Tree data structure, but the solution is not clear to me.

Here the Problem called Finding Magic Triplets, it can be found in Uva online judge:

(a + b^2) mod k = c^3 mod k, where a<=b<=c and 1 <= a, b, c <= n.

given n and k (1 <= n, k <= 10^5), how many different magic triplets exist for known values of n and k. A triplet is different from another if any of the three values is not same in both triplets.

and here the solution that I found:

#include <cstdio>
#include <cstring>
using namespace std;

typedef long long int64;

const int MAX_K = (int)(1e5);

int N, K;

struct BinaryIndexedTree{

    typedef int64 bit_t;

    static const int MAX_BIT = 3*MAX_K + 1;
    bit_t data[MAX_BIT+1];
    int SIZE;

    void init(int size){
        memset(data, 0, sizeof(data));
        SIZE = size;
    }

    bit_t sum(int n){
        bit_t ret = 0;
        for(;n;n-=n&-n){
            ret += data[n];
        }
        return ret;
    }

    bit_t sum(int from, int to){
        return sum(to)-sum(from);
    }

    void add(int n, bit_t x){
        for(n++;n<=SIZE;n+=n&-n){
            data[n]+=x;
        }
    }
};

BinaryIndexedTree bitree;


void init(){
    scanf("%d%d", &N, &K);
}

int64 solve(){
    bitree.init(2*K+1);

    int64 ans = 0;
    for(int64 i=N; i>=1; i--){
        int64 b = i * i % K, c = i * i * i % K;
        bitree.add(c, 1);
        bitree.add(c+K, 1);
        bitree.add(c+2*K, 1);
        int64 len = i;
        if(len >= K){
            ans += (len / K) * bitree.sum(K);
            len %= K;
        }
        if(len > 0){
            ans += bitree.sum(b + 1, b + len + 1);
        }
    }

    return ans;
}

int main(){
    int T;
    scanf("%d", &T);
    for(int i=0; i<T; i++){
        init();
        printf("Case %d: %lld\n", i+1, solve());
    }

    return 0;
}

Answer 1

Are you determined to use BITs? I would have thought ordinary arrays would do. I would start by creating three arrays of size k, where arrayA[i] = the number of values of a in range equal to i mod k, arrayB[i] = the number of values of b in range where b^2 = i mod k, and arrayC[i] = the number of values of c in range where c^3 = i mod k. N and k are both <= 10^5 so you could just consider each value of a in turn, b in turn, and c in turn, though you can be cleverer if k is much smaller than n, because will be some sort of fiddly fence-post counting expression that allows you to work out how many numbers in the range 0..n are equal to i mod k for each i.

Given those three arrays then consider each possible pair of numbers i, j where 0<=i,j < k and work out that there are arrayA[i] * arrayB[j] pairs which have those values mod k. Sum these up in arrayAB[i + j mod k] to find the number of ways that you can chose a + b^2 mod k = x for 0<=x < k. Now you have two arrays arrayAB and arrayC, where arrayAB[i] * arrayC[i] is the number of ways of finding a triple where a + b^2 = c^3] = i, so sum this over all 0<=i < k to get your answer.