Number of Ways Where Square of Number Is Equal to Product of Two Numbers LeetCode solution in Python

1
# Runtime: 753 ms (Top 50.00%) | Memory: 14 MB (Top 47.37%)
2
class Solution:
3
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
4
        sqr1, sqr2 = defaultdict(int), defaultdict(int)
5
        m, n = len(nums1), len(nums2)
6
        for i in range(m):
7
            sqr1[nums1[i] ** 2] += 1
8
        for j in range(n):
9
            sqr2[nums2[j] ** 2] += 1
10

11
        res = 0
12
        for i in range(m - 1):
13
            for j in range(i + 1, m):
14
                if nums1[i] * nums1[j] in sqr2:
15
                    res += sqr2[nums1[i] * nums1[j]]
16

17
        for i in range(n - 1):
18
            for j in range(i + 1, n):
19
                if nums2[i] * nums2[j] in sqr1:
20
                    res += sqr1[nums2[i] * nums2[j]]
21
        return res