Как я могу решить эту проблему с помощью BIT?

Я нашел хорошую математическую задачу, но до сих пор не могу ее решить. Я попытался найти одно решение с помощью Google и обнаружил, что его можно решить с помощью структуры данных двоичного индексированного дерева, но решение мне не ясно.

Вот Задача под названием «Поиск волшебных троек», ее можно найти в онлайн-судье Ува:

(a + b^2) mod k = c^3 mod k, где a‹=b‹=c и 1 ‹= a, b, c ‹= n.

при заданных n и k (1 ‹ = n, k ‹ = 10 ^ 5) сколько различных магических троек существует для известных значений n и k. Триплет отличается от другого, если какое-либо из трех значений не совпадает в обоих триплетах.

и вот решение, которое я нашел:

#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;
}