Calculate number of pairs a,b where a² - b² = n

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² - 1² = 8² - 7².

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¹³ so it was taking a very long time for larger values of n.

Since a² - b² = (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-digit 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 ?
Explanation

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 (

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.

Edit:

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^(-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

print(square_pairs_count(155235236))

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 t

StackExchange Code Review Q#139709, answer score: 7