Problem

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.

Examples

Example 1

1
2
3

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

Example 2

1
2
3
4
5

Input: nums1 = [1,1], nums2 = [1,1,1]
Output: 9
Explanation: All Triplets are valid, because 12 = 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

1
2
3
4
5

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].

Constraints

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

Solution

Method 1 – Hash Map for Pair Products

Intuition

For each number in one array, count how many pairs in the other array multiply to its square. Use a hash map to count all pair products efficiently.

Approach

  1. For each number in nums1, count pairs in nums2 whose product equals its square. Do the same for nums2 and nums1.
  2. For each array, build a hash map of product counts for all pairs (j < k).
  3. For each number in the other array, add the count of pairs whose product equals its square.
  4. Sum both directions for the answer.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class Solution {
public:
    int numTriplets(vector<int>& nums1, vector<int>& nums2) {
        return count(nums1, nums2) + count(nums2, nums1);
    }
    int count(vector<int>& a, vector<int>& b) {
        unordered_map<long long, int> prod;
        int n = b.size();
        for (int i = 0; i < n; ++i) {
            for (int j = i+1; j < n; ++j) {
                prod[1LL * b[i] * b[j]]++;
            }
        }
        int ans = 0;
        for (int x : a) ans += prod[1LL * x * x];
        return ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

func numTriplets(nums1, nums2 []int) int {
    return count(nums1, nums2) + count(nums2, nums1)
}
func count(a, b []int) int {
    prod := map[int64]int{}
    for i := 0; i < len(b); i++ {
        for j := i+1; j < len(b); j++ {
            prod[int64(b[i])*int64(b[j])]++
        }
    }
    ans := 0
    for _, x := range a {
        ans += prod[int64(x)*int64(x)]
    }
    return ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class Solution {
    public int numTriplets(int[] nums1, int[] nums2) {
        return count(nums1, nums2) + count(nums2, nums1);
    }
    private int count(int[] a, int[] b) {
        Map<Long, Integer> prod = new HashMap<>();
        int n = b.length;
        for (int i = 0; i < n; i++) {
            for (int j = i+1; j < n; j++) {
                long p = 1L * b[i] * b[j];
                prod.put(p, prod.getOrDefault(p, 0) + 1);
            }
        }
        int ans = 0;
        for (int x : a) ans += prod.getOrDefault(1L * x * x, 0);
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class Solution {
    fun numTriplets(nums1: IntArray, nums2: IntArray): Int {
        fun count(a: IntArray, b: IntArray): Int {
            val prod = mutableMapOf<Long, Int>()
            for (i in b.indices) for (j in i+1 until b.size)
                prod[b[i].toLong() * b[j].toLong()] = prod.getOrDefault(b[i].toLong() * b[j].toLong(), 0) + 1
            var ans = 0
            for (x in a) ans += prod.getOrDefault(x.toLong() * x.toLong(), 0)
            return ans
        }
        return count(nums1, nums2) + count(nums2, nums1)
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution:
    def numTriplets(self, nums1: list[int], nums2: list[int]) -> int:
        def count(a: list[int], b: list[int]) -> int:
            from collections import Counter
            prod = Counter()
            n = len(b)
            for i in range(n):
                for j in range(i+1, n):
                    prod[b[i]*b[j]] += 1
            ans = 0
            for x in a:
                ans += prod[x*x]
            return ans
        return count(nums1, nums2) + count(nums2, nums1)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

impl Solution {
    pub fn num_triplets(nums1: Vec<i32>, nums2: Vec<i32>) -> i32 {
        fn count(a: &Vec<i32>, b: &Vec<i32>) -> i32 {
            use std::collections::HashMap;
            let mut prod = HashMap::new();
            for i in 0..b.len() {
                for j in i+1..b.len() {
                    *prod.entry((b[i] as i64)*(b[j] as i64)).or_insert(0) += 1;
                }
            }
            let mut ans = 0;
            for &x in a {
                ans += *prod.get(&(x as i64 * x as i64)).unwrap_or(&0);
            }
            ans
        }
        count(&nums1, &nums2) + count(&nums2, &nums1)
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution {
    numTriplets(nums1: number[], nums2: number[]): number {
        function count(a: number[], b: number[]): number {
            const prod = new Map<number, number>()
            for (let i = 0; i < b.length; i++)
                for (let j = i+1; j < b.length; j++)
                    prod.set(b[i]*b[j], (prod.get(b[i]*b[j]) ?? 0) + 1)
            let ans = 0
            for (let x of a) ans += prod.get(x*x) ?? 0
            return ans
        }
        return count(nums1, nums2) + count(nums2, nums1)
    }
}

Complexity

  • ⏰ Time complexity: O(m^2 + n^2), where m and n are the lengths of nums1 and nums2. We count all pairs in each array.
  • 🧺 Space complexity: O(m^2 + n^2), for storing all pair products in hash maps.