# Counting triplets from three arrays satisfying the equation x^2 = yz

Let's say I have three arrays of positive integers X, Y and Z. You can assume that each of the arrays have atmost $$N$$ elements. All the arrays contain unique elements.

You're allowed to form a triplet of integers $$(x,y,z)$$ such that $$x$$, $$y$$ and $$z$$ are elements from arrays X, Y and Z respectively.

I want to count the number of such triplets satisfying the equation $$x^2 = yz$$. I can only think of an $$O(N^2)$$ solution using simple hashing. I was wondering if there exists a better solution with a time complexity of $$O(N \lg N)$$ at least, if not $$O(N)$$.

• How large are the integers? Can you share more about the context where you encountered this problem? Can you credit the original source of the problem, if you saw it somewhere?
– D.W.
Feb 11, 2020 at 20:49
• "Obviously", a non-concurrent $O(N\lg N)$ time solution will be in $O(N\lg N)$ space. Feb 11, 2020 at 21:02
• Worst case will be difficult. For random largish numbers, you could sort all arrays by largest prime factor, and if a number in Y has a prime factor p, then there must be a number in Z with a factor p, p^3 etc. and a corresponding number in X - this may allow you to remove lots of numbers, and reduce the number of cases a lot. Feb 11, 2020 at 22:35
• @D.W. I was trying to solve this problem after inspiring from the Two-sum Problem. It can be solved in $O(N)$. I was wondering if this problem (for a 3-variable equation) can be solved in $O(N \lg N)$, if not as good as $O(N)$. It's not a textbook problem so it doesn't really have any value limits as such. Let's say there's no restriction on the amount of memory that you can use, for simplicity. Feb 12, 2020 at 2:54
• If the integers were random and small, you could factor them, representing $x$ by the vector $e=(e_1,\dots,e_k)$ where $x=p_1^{e_1} \cdots p_k^{e_k}$ and $p_1,\dots,p_k$ is the list of small primes; then search for $y,z$ whose vectors $f,g$ satisfy $f \equiv g \pmod 2$ (this can be done in $O(n)$ time using a hashtable); for each such pair, such whether $yz \in X$. I don't know whether this would be competitive with other approaches, but it might be better for some specific distributions on $X,Y,Z$.
– D.W.
Feb 12, 2020 at 5:56