Medium

题目描述

给定两个整数数组 nums1 和 nums2，返回满足以下规则的三元组（类型1和类型2）的数量：

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

示例 1：

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

示例 2：

输入：nums1 = [1,1], nums2 = [1,1,1]
输出：9
解释：所有三元组都有效，因为 1² = 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)

示例 3：

输入：nums1 = [7,7,8,3], nums2 = [1,2,9,7]
输出：2
解释：有2个有效的三元组
类型1：(3,0,2)。nums1[3]² = nums2[0] * nums2[2]
类型2：(3,0,1)。nums2[3]² = nums1[0] * nums1[1]

提示：

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

解题思路

这道题要找满足平方等于两数乘积的三元组数量。我们可以用哈希表优化查找过程。

基本思路：

  1. 对于类型1，枚举 nums1[i]，计算 target = nums1[i]²，然后在 nums2 中找满足 nums2[j] * nums2[k] = target 的 (j,k) 对
  2. 对于类型2，枚举 nums2[i]，计算 target = nums2[i]²，然后在 nums1 中找满足 nums1[j] * nums1[k] = target 的 (j,k) 对

优化策略：

  • 使用哈希表统计每个数的出现次数，避免重复遍历
  • 对于找两数乘积等于目标值的情况，枚举第一个数，用哈希表查找第二个数
  • 注意处理相同数字的情况：如果两个数相同，组合数为 count * (count-1) / 2；如果不同，组合数为 count1 * count2

实现细节：

  • 预先计算两个数组中所有数字的频次
  • 分别处理类型1和类型2的情况
  • 避免整数溢出问题（平方可能很大）

代码实现

class Solution {
public:
    int numTriplets(vector<int>& nums1, vector<int>& nums2) {
        return helper(nums1, nums2) + helper(nums2, nums1);
    }
    
private:
    int helper(vector<int>& nums1, vector<int>& nums2) {
        unordered_map<int, int> freq;
        for (int num : nums2) {
            freq[num]++;
        }
        
        int count = 0;
        for (int num : nums1) {
            long long target = (long long)num * num;
            for (auto& [val, cnt] : freq) {
                if (target % val == 0) {
                    long long other = target / val;
                    if (other == val) {
                        count += cnt * (cnt - 1) / 2;
                    } else if (freq.count(other) && other > val) {
                        count += cnt * freq[other];
                    }
                }
            }
        }
        return count;
    }
};

class Solution:
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        def helper(arr1, arr2):
            freq = Counter(arr2)
            count = 0
            for num in arr1:
                target = num * num
                for val, cnt in freq.items():
                    if target % val == 0:
                        other = target // val
                        if other == val:
                            count += cnt * (cnt - 1) // 2
                        elif other in freq and other > val:
                            count += cnt * freq[other]
            return count
        
        return helper(nums1, nums2) + helper(nums2, nums1)

public class Solution {
    public int NumTriplets(int[] nums1, int[] nums2) {
        return Helper(nums1, nums2) + Helper(nums2, nums1);
    }
    
    private int Helper(int[] nums1, int[] nums2) {
        var freq = new Dictionary<int, int>();
        foreach (int num in nums2) {
            freq[num] = freq.GetValueOrDefault(num, 0) + 1;
        }
        
        int count = 0;
        foreach (int num in nums1) {
            long target = (long)num * num;
            foreach (var pair in freq) {
                int val = pair.Key;
                int cnt = pair.Value;
                if (target % val == 0) {
                    long other = target / val;
                    if (other == val) {
                        count += cnt * (cnt - 1) / 2;
                    } else if (freq.ContainsKey((int)other) && other > val) {
                        count += cnt * freq[(int)other];
                    }
                }
            }
        }
        return count;
    }
}

var numTriplets = function(nums1, nums2) {
    function countTriplets(arr1, arr2) {
        let count = 0;
        for (let i = 0; i < arr1.length; i++) {
            let target = arr1[i] * arr1[i];
            for (let j = 0; j < arr2.length; j++) {
                for (let k = j + 1; k < arr2.length; k++) {
                    if (arr2[j] * arr2[k] === target) {
                        count++;
                    }
                }
            }
        }
        return count;
    }
    
    return countTriplets(nums1, nums2) + countTriplets(nums2, nums1);
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度大小
时间复杂度O(m × k + n × k)，其中 m、n 分别是两个数组的长度，k 是不同数字的个数
空间复杂度O(k)，用于存储哈希表