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)
- if m < t and sq mod m is same as 0, then
- 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