In one site there was one question to make one program where you will be given a positive integer n and we have to calculate the number of such pairs (a, b),

where n = a² - b² and both a and b are positive integers.

for example 15 = 4^2 - 1^2 = 8^2 - 7^2.

Currently this question is closed and in my case I am still not able to do it properly. Still I won't be trying to find solution here but I just want to know that where I am doing wrong with my code.

At first I thought of simply increment a, b by 1 and check every equation. But time limit is 4 seconds and 1<= n<= 10^13 so it was taking very long time for larger values of n.

Since a^2 - b^2 = (a+b)(a-b) then I took one example 15 and solved for a = n and b = a-1 and observed that I got some special series:

(2n+1): n=1,2,3 ...   (for a = 2 and b = 1)

4n: n=2,3,4 ...       (for a = 3 and b = 1/b = 2)

6n+3: n = 2,3,4 ...   (for a = 4 and b = 1/b = 2/b = 3)

8n: n = 3,4,5 ...

Later I used this concept and made one program. It is working fine with 6 digits numbers(approx).

My program:

n = input()
count = 0
for i in range(1,n):
    first = (n-((2*i)-1))
    second = ((2*i)+2*(i-1))
    if first % second == 0 and first/second >= i:
        count += 1
    elif n%(4*i) == 0 and n/(4*i) >= (i+1):
        count +=1
print count

But for some larger values it is giving some error(16/30). If I make n in range to 100000 then I get few error as compared to previous one (27/30). But still it is taking time for larger input. So does my method is right or do I have to think of something else. If this is wrong then why ?


Let us say I want to find for no. 15 then according to formula (a+b)(a-b) = n total no. of possibilities :

(2+1)(2-1) = 3 # BLOCK 1

(3+1)(3-1) = 8 #BLOCK 2

(3+2)(3-2) = 5

(4+1)(4-1) = 15

(4+2)(4-2) = 12 #BLOCK 3

(4+3)(4-3) = 7

And so on ... till (8+7)(8-7) = 15

So if I go on like this then the last line of each block (in this case it is n = 3,5,7) we found that it is changing as (2n+1) where n = 1,2,3 .... (here we can think n as if from which block we are starting the count)

Now for every last second line (in this case it is 8,12) the equation will be 4n where n = 2,3,4 ....(since first block has only 1 line so there can't be any last second line so we will count from second block)

Repeat this process again and again .... We get:

2n+1, 4n, 6n+3, 8n, 10n+5, 12n, 14n+7, 16n ... and so on

If we simply try to find these equation every time for different lines then we find out that we can use this concept in program. If any no. 'n' is divisible by this equation it means it has occurred on those lines of block. Means each time it is divisible i will make count + 1.

  • \$\begingroup\$ Could you explain more of what you're trying to do in your solution? I didn't understand the "special series" in particular. For example, what's the significance of "6n+3" and "b = 1/b = 2/b = 3"? \$\endgroup\$
    – smarx
    Aug 26, 2016 at 16:46
  • \$\begingroup\$ @smarx I edited it please have a look \$\endgroup\$
    – Shashank
    Aug 26, 2016 at 17:35
  • \$\begingroup\$ I've added explanation about your idea with "blocks" to my answer below \$\endgroup\$ Aug 28, 2016 at 16:03

1 Answer 1


You were correct to try a^2 - b^2 = (a+b)(a-b) venue: take your n and break it into p and q so that p > q and p * q = n.

Usually there are several ways of doing it and almost each way defines a pair you need. a would be (p + q)/2 and b would be (p - q)/2. Thus, you only need pairs where either both p and q are even, or they are both odd.

You don't really need calculating a and b, or p and q, just number of pairs.

Thus, break your n into prime divisors and count distinct ways to divide those divisors into 2 different sets so that either both sets contain at least one 2, or neither do.


Ok, some more math:

Case with n=1 is trivial: no solutions. For all other cases we can assume that n has nonempty set of prime divisors.

Let's start with case when all prime divisors are distinct and 2 is not present.

Each divisor can be placed either in set 1 or in set 2. Since sets will be obviously different, we're guaranteed p > q and since we're working with n's divisors, we can be sure that p * q = n.

Divisors' placement can be described bijectively by a string of bits with the length as number of divisors. Hence, we'll have 2^n placements, 2^(<number of divisors>-1) pairs of p and q ( -1 because swapping sets gives same p and q).

Now what if some prime divisor (except for 2, which is still not present) is present k>1 times?

For all other divisors, situation is exactly as above, but our special divisor will provide *(k+1), instead of *2 to number of pairs. Well, actually, *2 of other divisors is just a special case with k = 1.

So, if 2 is not present, we go through n's prime factorisation and for each divisor and its respective k number of occurences, we multiply our number of placements by k+1. If result is even - there were odd ks, meaning that in no placement our 2 sets were equal. We still counted every pair of p and q twice (once as set1, set2 and once as set2, set1) so we divide k+1 product by 2 and that's our answer. If result is not even - we have even quantity of every divisor, meaning that n is a square. That means we counted twice every pair of p and q but we've also counted pair of sqrt(n) once - which we shouldn't have! We would have to substract 1 before dividing by 2.

If 2 is present, things don't get much complicated either. If there's only one 2 we return 0 - because there's no a and b that work. If there's more than one - we say that one goes to each set every time so we simply drop 2 twos and stop concerning ourselves with them, repeating steps above.

And now for some code:

from collections import Counter
from functools import reduce
from operator import mul

def product(stuff): return reduce(mul, stuff)

def square_pairs_count(n):
    """Return number of pairs (a,b) so that a*a - b*b = n

    (Algo explanation goes here)
    if n % 2 == 0 and n % 4 != 0:
        return 0 # (a+b) is odd and (a-b) is even (or vice versa)? No way!
    factorisation = Counter(prime_factorisation(n))
    if factorisation[2] > 1:
        # We use 2 '2's to guarantee that both (a+b) and (a-b) are even
        factorisation[2] -= 2
    factor_count_product = product(k + 1 for k in factorisation.values())

    if factor_count_product % 2 == 1:
        # n is a square and we counted root-root pair as (a+b) and (a-b), undo that
        factor_count_product -= 1

    return factor_count_product // 2 # we've counted every pair twice, so //2


Of course, this code could be called comprehensible only in context of all math above. In the wild, most of the code would be a docstring explaining what happens.

Number of divisors grows pretty slowly with growth of n. It all boils down to how fast you produce them prime factors.

Your idea with "blocks":

You can notice that "block" number c contains a = c + 1 and all bs that would give n > 0 (but not necessarily integer).

If you shift block numbering so that (2+1)(2-1) is actually block 2 then you will have block number equal to a used in this block with block 1 being empty. (because no b and n would satisfy our needs).

You could iterate through each pair like:

for a in range(2, max_a(n)):
    for b in range(1, a):
        # test our (a-b)(a+b)

However, that means iterating a lot.

You could notice that each block can produce no more than 1 solution, so we could iterate through blocks instead. However, we'd have to check if a block really produces a solution:

for a in range(2, max_a(n)):
    b_squared = a*a - n
    # check if b_squared is actually a square

The trick will be in determining max_a(n).

Surely, if 2*a - 1 > n, then a^2 - (a-1)^2 > n and thus for any b we're interested in, a^2 - b^2 > n.That makes our max_a something like (n + 1)//2, or n//2 + 1 for sake of corner cases.

That's too much! we'll have to check around 10^13 blocks, and that's a lot. Well, checking couple hundred thousand primes is not a simple thing either, but it's way faster.

So, unless you can lower the boundary for a, the code would still take lots of time.

Modulo sieving

If you are interested in generative solutions, I'd recommend picking some number N, rewriting equation modulo N and looking for moduli N for a and b that are possible with given n. If, for example, N is around hundred buth there's only 10 possible modulo pairs that work - you would have to iterate only through 10 (a,b) pairs out of 10 000.

For example, if N = 3 and n mod 3 = 2 then a mod 3 = 0 and either b mod 3 = 1 or b mod 3 = 2. Of each 9 consecutive possible options for (a,b) only 2 remain. Of each 3 consecutive possible options for a only 1 remains.

For example, if N = 5 and n mod 5 = 1 then only a^2 mod 5 = 1, b^2 mod 5 = 0 and a^2 mod 5 = 0, b^2 mod 5 = 4. Of 25 possible options for (a,b) pair this leaves only 4. Of 5 possible options for a, it leaves only 3.

Moreover, if you have calculated allowed modulo pairs for some M and some N you can then calculate allowed modulo pairs for M*N and instead of iterating possible pairs with steps of size N (or M), go with M*N steps.

Combining 2 examples above, you'll get that you only need to check 3 out of 15 options for a.

Using these 2 tricks (modulo sieving and combining) you can iterate with ever-increasing steps without missing anything. However, problem does not ask for generating those pairs, only for counting their number, so that's a bit excessive.

If you find that option interesting, take a closer look at sieves and especially wheels part.

  • \$\begingroup\$ You can actually drop if factor_count_product % 2 == 1: factor_count_product -= 1. For every p×q, there will also appear q×p, so factor_count_product is always even, except when n is a perfect square — in which case √n×√n appears just once, making factor_count_product odd. The // 2 will take care of that by truncating down. \$\endgroup\$ Aug 27, 2016 at 22:20
  • 1
    \$\begingroup\$ Instead of if factorisation[2] == 1: return 0, consider checking if n % 2 == 0 and n % 4 != 0: return 0 before doing prime_factorisation(n). It's a quick check that could potentially save a lot of calculation. \$\endgroup\$ Aug 27, 2016 at 22:24
  • \$\begingroup\$ > // 2 will take care of that by truncating down I think accounting for that case explicitly is better \$\endgroup\$ Aug 28, 2016 at 10:18

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.