p251 get square divisors

p251_square_divisors.py

def gen_multiples(primes, multiplicities, limit):
    if len(primes) == 0:
        return []
    p = 1
    multiples = []
    for i in range(multiplicities[0] + 1):
        if p > 1:
            multiples.append(p)
            
        if len(primes) > 1:
            for x in gen_multiples(primes[1:], multiplicities[1:], limit//p):
                if p*x < limit:
                    multiples.append(p*x)
        
        p *= primes[0]
        
    return multiples

def get_sq_divisors(a, limit):
    pf1 = prime_factors((a + 1)//3)
    pf2 = prime_factors((8*a - 1)//3)
    
    pf = {}
    for p in pf1:
        pf[p] = 2*pf1[p]
    for p in pf2:
        pf[p] = pf2[p] + pf.get(p, 0)
    
    primes = []
    multiplicities = []
    for p, e in pf.items():
        mult = e//2
        if mult > 0:
            primes.append(p)
            multiplicities.append(mult)
    
    return [x for x in gen_multiples(primes, multiplicities, limit)]