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 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).

To solve this, we will follow these steps:

• cnt1 := a map to hold each elements and its count of nums1
• cnt2 := a map to hold each elements and its count of nums2
• Define a function triplets() . This will take arr1, arr2
• ans := 0
• for each t, v in items() of arr1, do
  • k := arr2[t] if it is there, otherwise 0
  • tmp := k*(k - 1) / 2
  • sq := t * t
  • for each m in arr2, do
    • if m < t and sq mod m is same as 0, then
      • tmp := tmp + (arr2[m] if it is there, otherwise 0) * (arr2[quotient of sq/m] if it is there, otherwise 0)
• ans := ans + tmp * v
• return ans
• From the main method, do the following−
• return triplets(cnt1, cnt2) + triplets(cnt2, cnt1)

Example

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)
nums1 = [7,4]
nums2 = [5,2,8,9]
print(solve(nums1, nums2))

Input

[7,4],[5,2,8,9]

Output

2