Formatted question description: https://leetcode.ca/all/1577.html

1577. Number of Ways Where Square of Number Is Equal to Product of Two Numbers (Medium)

Given two arrays of integers nums1 and nums2, return the number of triplets formed (type 1 and type 2) under the following rules:

• Type 1: Triplet (i, j, k) if nums1[i]2 == nums2[j] * nums2[k] where 0 <= i < nums1.length and 0 <= j < k < nums2.length.
• Type 2: Triplet (i, j, k) if nums2[i]2 == nums1[j] * nums1[k] where 0 <= i < nums2.length and 0 <= j < k < nums1.length.

Example 1:

Input: nums1 = [7,4], nums2 = [5,2,8,9]
Output: 1
Explanation: Type 1: (1,1,2), nums1[1]^2 = nums2[1] * nums2[2]. (4^2 = 2 * 8).


Example 2:

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

Related Topics:
Hash Table, Math

Solution 1.

Type 1 and Type 2 are symmetrical so we can define a function count(A, B) which returns the count of the Type 1 triples. The answer is count(A, B) + count(B, A).

For count(A, B), we can use a unordered_map<int, int> m to store the frequency of the numbers in B. Then for each number a in A, the target value is a * a. We traverse the map m to count the triplets.

For each entry (b, cnt) in m:

• If target is not divisible by b or target / b is not in m, there is no triplets, skip.
• If target / b == b, we need to pick 2 out of cnt numbers so we can add cnt * (cnt - 1) triplets to the answer.
• Otherwise, we can add cnt * m[target / b] triplets to the answer.

Since we count the the pairs in B twice, we need to divide the answer by 2 before returning.

// OJ: https://leetcode.com/problems/number-of-ways-where-square-of-number-is-equal-to-product-of-two-numbers/
// Time: O(N^2)
// Space: O(N)
class Solution {
int count(vector<int> &A, vector<int> &B) {
int ans = 0;
unordered_map<int, int> m;
for (int n : B) m[n]++;
for (int a : A) {
long target = (long)a * a;
for (auto &[b, cnt] : m) {
if (target % b || m.count(target / b) == 0) continue;
if (target / b == b) ans += cnt * (cnt - 1);
else ans += cnt * m[target / b];
}
}
return ans / 2;
}
public:
int numTriplets(vector<int>& A, vector<int>& B) {
return count(A, B) + count(B, A);
}
};

• class Solution {
public int numTriplets(int[] nums1, int[] nums2) {
Map<Integer, Integer> map1 = new HashMap<Integer, Integer>();
Map<Integer, Integer> map2 = new HashMap<Integer, Integer>();
for (int num : nums1) {
int count = map1.getOrDefault(num, 0) + 1;
map1.put(num, count);
}
for (int num : nums2) {
int count = map2.getOrDefault(num, 0) + 1;
map2.put(num, count);
}
Set<Integer> set1 = map1.keySet();
Set<Integer> set2 = map2.keySet();
int triplets = 0;
for (int num1 : set1) {
int count1 = map1.get(num1);
long square = (long) num1 * num1;
for (int num2 : set2) {
if (square % num2 == 0) {
int num3 = (int) (square / num2);
if (num2 == num3) {
int count2 = map2.get(num2);
int curTriplets = count1 * count2 * (count2 - 1) / 2;
triplets += curTriplets;
} else if (num2 < num3 && set2.contains(num3)) {
int count2 = map2.get(num2), count3 = map2.get(num3);
int curTriplets = count1 * count2 * count3;
triplets += curTriplets;
}
}
}
}
for (int num1 : set2) {
int count1 = map2.get(num1);
long square = (long) num1 * num1;
for (int num2 : set1) {
if (square % num2 == 0) {
int num3 = (int) (square / num2);
if (num2 == num3) {
int count2 = map1.get(num2);
int curTriplets = count1 * count2 * (count2 - 1) / 2;
triplets += curTriplets;
} else if (num2 < num3 && set1.contains(num3)) {
int count2 = map1.get(num2), count3 = map1.get(num3);
int curTriplets = count1 * count2 * count3;
triplets += curTriplets;
}
}
}
}
return triplets;
}
}

############

class Solution {
public int numTriplets(int[] nums1, int[] nums2) {
Map<Integer, Integer> cnt1 = new HashMap<>();
Map<Integer, Integer> cnt2 = new HashMap<>();
for (int v : nums1) {
cnt1.put(v, cnt1.getOrDefault(v, 0) + 1);
}
for (int v : nums2) {
cnt2.put(v, cnt2.getOrDefault(v, 0) + 1);
}
long ans = 0;
for (var e1 : cnt1.entrySet()) {
long a = e1.getKey(), x = e1.getValue();
for (var e2 : cnt2.entrySet()) {
long b = e2.getKey(), y = e2.getValue();
if ((a * a) % b == 0) {
long c = a * a / b;
if (b == c) {
ans += x * y * (y - 1);
} else {
ans += x * y * cnt2.getOrDefault((int) c, 0);
}
}
if ((b * b) % a == 0) {
long c = b * b / a;
if (a == c) {
ans += x * (x - 1) * y;
} else {
ans += x * y * cnt1.getOrDefault((int) c, 0);
}
}
}
}
return (int) (ans >> 1);
}
}

• // OJ: https://leetcode.com/problems/number-of-ways-where-square-of-number-is-equal-to-product-of-two-numbers/
// Time: O(N^2)
// Space: O(N)
class Solution {
int count(vector<int> &A, vector<int> &B) {
int ans = 0;
unordered_map<int, int> m;
for (int n : B) m[n]++;
for (int a : A) {
long target = (long)a * a;
for (auto &[b, cnt] : m) {
if (target % b || m.count(target / b) == 0) continue;
if (target / b == b) ans += cnt * (cnt - 1);
else ans += cnt * m[target / b];
}
}
return ans / 2;
}
public:
int numTriplets(vector<int>& A, vector<int>& B) {
return count(A, B) + count(B, A);
}
};

• class Solution:
def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
cnt1 = Counter(nums1)
cnt2 = Counter(nums2)
ans = 0
for a, x in cnt1.items():
for b, y in cnt2.items():
if a * a % b == 0:
c = a * a // b
if b == c:
ans += x * y * (y - 1)
else:
ans += x * y * cnt2[c]
if b * b % a == 0:
c = b * b // a
if a == c:
ans += x * (x - 1) * y
else:
ans += x * y * cnt1[c]
return ans >> 1


• func numTriplets(nums1 []int, nums2 []int) (ans int) {
cnt1 := map[int]int{}
cnt2 := map[int]int{}
for _, v := range nums1 {
cnt1[v]++
}
for _, v := range nums2 {
cnt2[v]++
}
for a, x := range cnt1 {
for b, y := range cnt2 {
if a*a%b == 0 {
c := a * a / b
if b == c {
ans += x * y * (y - 1)
} else {
ans += x * y * cnt2[c]
}
}
if b*b%a == 0 {
c := b * b / a
if a == c {
ans += x * (x - 1) * y
} else {
ans += x * y * cnt1[c]
}
}
}
}
ans /= 2
return
}