Input: nums1 = [1,1], nums2 = [1,1,1]
Output: 9
Explanation: All Triplets are valid, because 1^2 = 1 * 1.
Type 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].
Type 2: (0,0,1), (1,0,1), (2,0,1). nums2[i]^2 = nums1[j] * nums1[k].
Example 3:
Input: nums1 = [7,7,8,3], nums2 = [1,2,9,7]
Output: 2
Explanation: There are 2 valid triplets.
Type 1: (3,0,2). nums1[3]^2 = nums2[0] * nums2[2].
Type 2: (3,0,1). nums2[3]^2 = nums1[0] * nums1[1].
Example 4:
Input: nums1 = [4,7,9,11,23], nums2 = [3,5,1024,12,18]
Output: 0
Explanation: There are no valid triplets.
Constraints:
1 <= nums1.length, nums2.length <= 1000
1 <= nums1[i], nums2[i] <= 10^5
Solution:
class Solution {
public int numTriplets(int[] nums1, int[] nums2) {
int res = 0;
for (int val : nums1) {
long v = (long) val;
res += twoSum(v * v, nums2);
}
for (int val : nums2) {
long v = (long) val;
res += twoSum(v * v, nums1);
}
return res;
}
private int twoSum(long target, int[] arr) {
Map<Long, Integer> map = new HashMap();
int count = 0;
for (int i = 0; i < arr.length; i ++) {
long curr = (long) arr[i];
long d = target / curr;
if (d * curr != target) continue;
if (map.containsKey(d)) {
count += map.get(d);
}
map.put(curr, map.getOrDefault(curr, 0) + 1);
}
return count;
}
}