PROBLEM LINKS

Practice
Contest

DIFFICULTY

EASY

PREREQUISITES

Simple Math, Dynamic Programming, Repeated Squaring

PROBLEM

You are given three integers d, m and N in the following relation

xd + yd + zd≡ m, modulo N

Find the number of solutions of the above equation where

0 ≤ x, y, z ≤ U

EXPLANATION

The first observation to be made is that the value of N is very small. We already know that

ab = (a % N)b, modulo N

Despite the fact that U is a large number, it is only relevant for us to find the solutions (X, Y, Z) which satisfy

0 ≤ X, Y, Z < N

and for each solution (X, Y, Z), calculate the number of values in the range

0 ≤ x, y, z ≤ U

such that

• x = X, modulo N
• y = Y, modulo N
• z = Z, modulo N

Let us say that

• A[t] = t^d % N
• 0 ≤ t < N

Now, we can iterate through the list of unique solutions (X, Y, Z)

for X = 0 to N-1
    for Y = 0 to N-1
        for Z = 0 to N-1
            if (A[X] + A[Y] + A[Z]) % N = m
                process(X, Y, Z)

The complexity of calculating A[] is O(N log d) if we use repeated squaring.

In the worst case, the number of times process() is called is O(N³).

For each solution (X, Y, Z) we wish add the solutions (x, y, z) from the complete range. The number of values of x, satisfying

• x = X, modulo N
• 0 ≤ x ≤ U

is

number-of-sols(X) =

floor(U / N) + {
    1, if (U % N > 0 AND X ≤ U % N),
    0, otherwise
}

Now, the number of solutions (x, y, z) for (X, Y, Z) is simply

answer = 0

process(X, Y, Z) =
    answer = answer +
        number-of-sols(X) *
        number-of-sols(Y) *
        number-of-sols(Z)

Thus, the answer can be found in O(N³) time.

CODING COMMENTARY

The answer has to be found modulo 10^9</sup+7. This leaves a lot of scope for integer overflow. Make sure that your implementation uses 64-bit integers for all intermediate results, at the least.

64-bit integers can handle the overflow of multiplying two 32-bit integers, but not three. You must find the modular result after each product.

SETTER'S SOLUTION

Can be found here.

TESTER'S SOLUTION

Can be found here.