题目描述

给定两个整数数组 nums1 和 nums2，请你返回根据以下规则形成的三元组的数目（类型 1 和类型 2 ）：

• 类型 1：三元组 (i, j, k)，如果 nums1[i]^2 == nums2[j] * nums2[k] 其中 0 <= i < nums1.length 且 0 <= j < k < nums2.length
• 类型 2：三元组 (i, j, k)，如果 nums2[i]^2 == nums1[j] * nums1[k] 其中 0 <= i < nums2.length 且 0 <= j < k < nums1.length

样例

输入：nums1 = [7,4], nums2 = [5,2,8,9]
输出：1
解释：类型 1：(1,1,2), nums1[1]^2 = nums2[1] * nums2[2] (4^2 = 2 * 8)

输入：nums1 = [1,1], nums2 = [1,1,1]
输出：9
解释：所有三元组都符合题目要求，因为 1^2 = 1 * 1
类型 1：(0,0,1), (0,0,2), (0,1,2), (1,0,1), (1,0,2), (1,1,2)
类型 2：(0,0,1), (1,0,1), (2,0,1)

输入：nums1 = [7,7,8,3], nums2 = [1,2,9,7]
输出：2
解释：有两个符合题目要求的三元组：
类型 1：(3,0,2), nums1[3]^2 = nums2[0] * nums2[2]
类型 2：(3,0,1), nums2[3]^2 = nums1[0] * nums1[1]

输入：nums1 = [4,7,9,11,23], nums2 = [3,5,1024,12,18]
输出：0
解释：不存在符合题目要求的三元组。

限制

• 1 <= nums1.length, nums2.length <= 1000
• 1 <= nums1[i], nums2[i] <= 10^5

算法

(哈希表) $O(n^2 + m^2)$

1. 用两个哈希表分别统计每个数组中，每个数字的平方出现的次数。
2. 暴力枚举 nums1 中的两个元素，求乘积后，答案累加哈希表 2 中的值；对于 nums2 同理。

时间复杂度

• 维护哈希表的时间复杂度为 $O(n+m)$。
• 暴力枚举判断的时间复杂度为 $O(n^2 + m^2)$，故总时间复杂度为 $O(n^2 + m^2)$。

空间复杂度

• 需要 $O(n + m)$ 的额外空间存储哈希表。

Go 代码

func numTriplets(nums1 []int, nums2 []int) int {
    seen1 := make(map[int64]int)
    seen2 := make(map[int64]int)

    for _, v := range nums1 {
        seen1[int64(v) * int64(v)]++
    }
    for _, v := range nums2 {
        seen2[int64(v) * int64(v)]++
    }

    n, m := len(nums1), len(nums2)
    var ans int

    for i := 0; i < n; i++ {
        for j := i+1; j < n; j++ {
            ans += seen2[int64(nums1[i]) * int64(nums1[j])]
        }
    }

    for i := 0; i < m; i++ {
        for j := i+1; j < m; j++ {
            ans += seen1[int64(nums2[i]) * int64(nums2[j])]
        }
    }

    return ans
}