Leetcode 1577. Number of Ways Where Square of Number Is Equal to Product of Two Numbers

文章作者：Tyan
博客：noahsnail.com  |  CSDN  |  简书

1. Description

2. Solution

解析： Version 1，分别计算两个数组的平方和以及所有组合乘积并统计对应值的个数，遍历每个数组平方和的个数，找到另一个数组对应的积的个数，二者相乘，加到三元组总个数中。Version 2进行进一步优化。

• Version 1

class Solution:
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        count = 0
        square1 = collections.defaultdict(int)
        square2 = collections.defaultdict(int)
        product1 = collections.defaultdict(int)
        product2 = collections.defaultdict(int)
        for x in nums1:
            temp = x ** 2
            square1[temp] = square1[temp] + 1
        for x in nums2:
            temp = x ** 2
            square2[temp] = square2[temp] + 1
        for i in range(len(nums1)):
            for j in range(i+1, len(nums1)):
                temp = nums1[i] * nums1[j]
                product1[temp] = product1[temp] + 1
        for i in range(len(nums2)):
            for j in range(i+1, len(nums2)):
                temp = nums2[i] * nums2[j]
                product2[temp] = product2[temp] + 1
        for k, v in square1.items():
            count += v * product2[k]
        for k, v in square2.items():
            count += v * product1[k]
        return count

• Version 2

class Solution:
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        count = 0
        product1 = collections.defaultdict(int)
        product2 = collections.defaultdict(int)
        for i in range(len(nums1)):
            for j in range(i+1, len(nums1)):
                temp = nums1[i] * nums1[j]
                product1[temp] = product1[temp] + 1
        for i in range(len(nums2)):
            for j in range(i+1, len(nums2)):
                temp = nums2[i] * nums2[j]
                product2[temp] = product2[temp] + 1
        for x in nums1:
            temp = x ** 2
            count += product2[temp]
        for x in nums2:
            temp = x ** 2
            square2[temp] = square2[temp] + 1
            count += product1[temp]
        return count

Reference

1. https://leetcode.com/problems/number-of-ways-where-square-of-number-is-equal-to-product-of-two-numbers/