LeetCode 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
解释：不存在符合题目要求的三元组

提示：

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

先对两个数组进行排序，然后相向双指针计算即可：

class Solution {
public:
    int numTriplets(vector<int>& nums1, vector<int>& nums2) {
        sort(nums1.begin(), nums1.end());
        sort(nums2.begin(), nums2.end());
        
        return getNum(nums1, nums2) + getNum(nums2, nums1);
    }

    int getNum(vector<int> &nums1, vector<int> &nums2) {
        int ans = 0;

        for (int i = 0; i < nums1.size(); ++i) {
            int left = 0;
            int right = nums2.size() - 1;

            while (left < right) {
                if ((long long)nums2[left] * nums2[right] < (long long)nums1[i] * nums1[i]) {
                    ++left;
                } else if ((long long)nums2[left] * nums2[right] > (long long)nums1[i] * nums1[i]) {
                    --right;
                } else {
                    // 如果left和right之间全是相同的数，此时任意两对数字乘积都合法
                    if (nums2[left] == nums2[right]) {
                        ans += (right - left + 1) * (right - left) / 2;
                        break;
                    }

					// 计算左边相同数字的个数
                    int leftSame = 1;
                    ++left;
                    while (nums2[left] == nums2[left - 1]) {
                        ++left;
                        ++leftSame;
                    }

					// 计算右边相同数字的个数
                    int rightSame = 1;
                    --right;
                    while (nums2[right] == nums2[right + 1]) {
                        --right;
                        ++rightSame;
                    }

                    ans += leftSame * rightSame;
                }
            }
        }

        return ans;
    }
};

如果nums1的长度为n，nums2的长度为m，则此算法时间复杂度为O(max(nlogn, mlogm, mn))，空间复杂度为O(max(logn, logm))。