Compute pairs from two lists that when multiplied make a perfect square

Question

You are given two lists and you have to find out the pairs which make a perfect square.

For example in:

a = [2, 6, 10, 13, 17, 18]

b = [3, 7, 8, 9, 11, 15]

There are two pairs (2,8) and (8,18).

Is there any method efficient than the brute-force? Here is my code which has a time complexity of O(n*m) (where n is the length of a and m is the length of b).

pl = []
a = [ 2, 6, 10, 13, 17,18]
b = [ 3, 7, 8, 9, 11, 15 ]
i = 0
while(i < len(a)):
    j = 0
    while(j < len(b)):
        p = a[i]*b[j]
        n = p**0.5
        u = int(n)
        if n == u:
            pl.append((a[i],b[j]))

        j = j+1


    i = i+1

print(pl)

This question has been asked before using C# here, but I don't get what they mean by "All we would need to store for each number is which of its prime factors have an odd count", so I can't implement this in my Python code.

Can someone explain to me how we might implement an efficient solution in Python?

Answer 1

The logic described in the linked question is as follows. A perfect square's prime factors will always be in pairs. For example, 36 = 2*2*3*3. We have two 3s and two 2s. Therefore, if we take any two numbers that multiply to form a perfect square, if we take the combined counts of each of their prime factors, we will also get even counts.

For example, 8 = 2*2*2, and 18 = 2*3*3. Combined, we get four 2s and two 3s.

Below is some Python code that implements this algorithm, using collections.Counter and sets to remove duplicates. First, we precompute all prime factorizations for each unique element in a and b to avoid redundancy, and store this in a dictionary. Then, we loop over pairs of elements from a and b, using itertools.product, and combine the prime factor counts. If every count is even, we add the sorted pair to a set.

Code:

from collections import Counter
import itertools

def prime_factors(n):
    """
    https://stackoverflow.com/a/22808285/12366110
    """
    i = 2
    factors = []
    while i * i <= n:
        if n % i:
            i += 1
        else:
            n //= i
            factors.append(i)
    if n > 1:
        factors.append(n)
    return Counter(factors)

a = [2, 6, 10, 13, 17,18]

b = [3, 7, 8, 9, 11, 15]

prime_factors = {i: prime_factors(i) for i in set(a) | set(b)}

rv = set()

for a, b in itertools.product(a, b):
    combined_counts = prime_factors[a] + prime_factors[b]
    if all(v%2 == 0 for v in combined_counts.values()):
        rv.add(tuple(sorted([a, b])))

Output:

>>> rv
{(2, 8), (8, 18)}