I have a integers from 1 to n. I randomly allot every integer into one of three sets A, B and C (A ∩ B = B ∩ C = C ∩ A = Ø). Every integer does belong to one set. So I need to calculate all combination of elements (a,b) such that a ∈ A, b ∈ B, and the geometric mean of a,b belongs to C. Basically sqrt(a*b) ∈ C.
My solution is to first mark on an array of size n whether every element went into set A,B or C. Then I loop through the array for all elements that belong to A. When I encounter one, I again loop through for all elements that belong to B. If array[sqrt(a*b)] == C, then I add (a, b, sqrt(a,b)) as one possible combination. Then I do the same for the entire array, which is O(n^2).
Is there a more optimal solution possible?
It can be done with better complexity than O(n^2). The solution sketched here is in O(n * sqrt(n) * log(n)).
The main idea is the following:
One critical part in the idea above is decomposing a into s * t^2, i.e. finding the primes that have odd power in a's factorization.
This can be done using a pre-processing step, that finds the prime factors (but not also their powers) of every number in {1, 2, .. n}, using a slightly modified sieve of Eratosthenes. This modified version would not only mark a number as "not prime" when iterating over the multiples of a prime, but would also append the current prime number to the list of the factors of the current multiple. The time complexity of this pre-processing step is n * sum{for each prime p < n}(1/p) = n * log(log(n)) -- see this for details.
Using the result of the pre-processing, which is the list of primes which divide a, we can find those primes with odd power in O(log(n)). This is achieved by dividing a by each prime in the list until it is no more divisible by that prime. If we made an odd number of divisions, then we use the current prime in s. After all divisions are done, the result will be equal to 1. The complexity of this is O(log(n)) because in the worst case we always divide the initial number by 2 (the smallest prime number), thus it will take at most log2(a) steps to reach value 1.
The complexity of the main step dominates the complexity of the preprocessing, thus the overall complexity of this approach is O(n * sqrt(n) * log(n)).
Remark: in the decomposition a = s * t^2, s is the product of the prime numbers in a with odd exponents, but their exponent is not used in s (i.e. s is just the product of those primes, with exponent 1). Only in this situation it is guaranteed that b should be of the form s * k^2. Indeed, since a * b = c * c, the prime factorization of the right hand side uses only even exponents, thus all primes from s should also appear in b with odd exponents, and all other primes from b's factorization should have even exponents.
Expanding on the following line: "we can restrict our search for the number b to those of the form b = s * k^2, but there are only O(sqrt(n)) numbers like this smaller than n".
Let's consider an example. Imagine that we have something like n = 10,000 and we are currently looking for solutions having a = 360 = 2^3 * 3^2 * 5. The primes with odd exponent in a's factorization are 2 and 5 (thus s = 2 * 5; a = 10 * 6^2).
Since a * b is a perfect square, it means that all primes in the prime factorization of a * b have even exponents . This implies that those two primes (2 and 5) need to also appear in b's factorization with odd exponents, and the rest of the exponents in b's prime factorization need to be even. Thus b is of the form s * k^2 = 10 * k ^ 2.
So we proved that b = 10 * k ^ 2. This is helpful, because we can now enumerate all the b values of this form quickly (in O(sqrt(n)). We only need to consider k = 1, k = 2, ..., k = (int)sqrt(n / 10). Larger values of k result in values of b larger than n. Each of these k values determines one b value, which we need to verify. Note that when verifying one of these b values, it should be first checked whether it indeed is in set B, which can be done in O(1), and whether sqrt(a * b) is in the set C, which can also be done in O(1).