1577. 数的平方等于两数乘积的方法数

题目

给你两个整数数组 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

示例 1：

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

示例 2：

输入：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), nums1[i]^2 = nums2[j] * nums2[k]
类型 2：(0,0,1), (1,0,1), (2,0,1), nums2[i]^2 = nums1[j] * nums1[k]

示例 3：

输入：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]

示例 4：

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

思路

这道题其实就是将类型一和二分成两个小问题，然后依次求解，唯一的难点在于两个指针的乘积与第一个相同时，要怎么统计符合条件的三元组。

• 当两个指针的值相同，由于数组是排序过的，所以中间的值都相同，所以可以说是随意排列组合，根据组合知识，用组合公式求解
• 当左指针后有m个相同值，右指针前有n个相同值，一共有m*n种组合

class Solution {
    public int numTriplets(int[] nums1, int[] nums2) {
        Arrays.sort(nums1);
        Arrays.sort(nums2);
        int count = 0;
        // 计算类型一
        for (int i = 0; i < nums1.length; i++) {
            count += countTriple((long) nums1[i] * nums1[i], nums2);
        }
        // 计算类型二
        for (int i = 0; i < nums2.length; i++) {
            count += countTriple((long) nums2[i] * nums2[i], nums1);
        }
        return count;
    }

    private int countTriple(long pow, int[] nums) {
        int count = 0;
        int j = 0, k = nums.length - 1;
        while (j < k) {
            long x = (long) nums[j] * nums[k];
            if (pow < x) {
                k--;
            } else if (pow > x) {
                j++;
            } else {
                // 因为是排过序的，所以如果两个指针的值相同，中间所有值都是相同的
                // 利用组合公式求解
                if (nums[j] == nums[k]) {
                    count += (k - j + 1) * (k - j) / 2;
                    break;
                }
                int m = 1, n = 1;
                while (j < k && nums[j] == nums[j + 1]) {
                    m++;
                    j++;
                }
                while (j < k && nums[k] == nums[k - 1]) {
                    n++;
                    k--;
                }
                count += m * n;
                j++;
                k--;
            }
        }
        return count;
    }
}