Code Implementation

class Solution:
    def numTriplets(self, nums1, nums2):
        from collections import Counter

        def countPairs(arr1, arr2):
            count2 = Counter(arr2)
            res = 0
            for a in arr1:
                target = a * a
                for b in count2:
                    if target % b == 0:
                        c = target // b
                        if c in count2:
                            if b == c:
                                res += count2[b] * (count2[b] - 1) // 2
                            elif b < c:
                                res += count2[b] * count2[c]
            return res

        return countPairs(nums1, nums2) + countPairs(nums2, nums1)
      

class Solution {
public:
    int numTriplets(vector<int>& nums1, vector<int>& nums2) {
        return countPairs(nums1, nums2) + countPairs(nums2, nums1);
    }
    
    int countPairs(vector<int>& arr1, vector<int>& arr2) {
        unordered_map<long long, int> freq;
        for (int n : arr2) freq[n]++;
        int res = 0;
        for (int a : arr1) {
            long long target = 1LL * a * a;
            for (auto& [b, cnt_b] : freq) {
                if (target % b == 0) {
                    long long c = target / b;
                    if (freq.count(c)) {
                        if (b == c)
                            res += cnt_b * (cnt_b - 1) / 2;
                        else if (b < c)
                            res += cnt_b * freq[c];
                    }
                }
            }
        }
        return res;
    }
};
      

class Solution {
    public int numTriplets(int[] nums1, int[] nums2) {
        return countPairs(nums1, nums2) + countPairs(nums2, nums1);
    }
    
    private int countPairs(int[] arr1, int[] arr2) {
        Map<Long, Integer> freq = new HashMap<>();
        for (int n : arr2) freq.put((long)n, freq.getOrDefault((long)n, 0) + 1);
        int res = 0;
        for (int a : arr1) {
            long target = 1L * a * a;
            for (long b : freq.keySet()) {
                if (target % b == 0) {
                    long c = target / b;
                    if (freq.containsKey(c)) {
                        if (b == c)
                            res += freq.get(b) * (freq.get(b) - 1) / 2;
                        else if (b < c)
                            res += freq.get(b) * freq.get(c);
                    }
                }
            }
        }
        return res;
    }
}
      

var numTriplets = function(nums1, nums2) {
    function countPairs(arr1, arr2) {
        let freq = new Map();
        for (let n of arr2) freq.set(n, (freq.get(n) || 0) + 1);
        let res = 0;
        for (let a of arr1) {
            let target = a * a;
            for (let [b, cnt_b] of freq.entries()) {
                if (target % b === 0) {
                    let c = target / b;
                    if (freq.has(c)) {
                        if (b === c)
                            res += cnt_b * (cnt_b - 1) / 2;
                        else if (b < c)
                            res += cnt_b * freq.get(c);
                    }
                }
            }
        }
        return res;
    }
    return countPairs(nums1, nums2) + countPairs(nums2, nums1);
};
      

Problem Description

Given two integer arrays nums1 and nums2, your task is to count the number of triplets where:

• Pick one element a from nums1, and two elements b and c from nums2 (they can be the same or different indices), such that a * a = b * c.
• Or, pick one element a from nums2, and two elements b and c from nums1, such that a * a = b * c.

The order of b and c does not matter, but distinct pairs of indices are counted separately. You cannot reuse the same element for both b and c unless they are at different indices.

Thought Process

At first glance, the problem looks like a brute-force candidate: for every element in one array, try all possible pairs in the other array and check if the equation holds. However, this would be too slow for large arrays, since the number of pairs grows quadratically.

To optimize, we recognize that for a given a, we want to count the number of pairs (b, c) such that b * c = a * a. This is a classic case for using a hash map to count frequencies, so we can efficiently check for the required products.

We also need to be careful to avoid counting duplicate pairs when b and c are the same, and to count each unordered pair only once.

Solution Approach

• Step 1: For each array, build a frequency map (hash map) that counts how many times each value appears.
• Step 2: For each element a in the first array, compute target = a * a.
• Step 3: For each unique number b in the second array, check if target % b == 0 (to ensure c is an integer).
• Step 4: If so, set c = target / b. If c is also in the frequency map, count the pairs:
  • If b == c: The number of pairs is freq[b] * (freq[b] - 1) / 2 (choose 2 from freq[b]).
  • If b < c: The number of pairs is freq[b] * freq[c].
  • We only consider b ≤ c to avoid double-counting unordered pairs.
• Step 5: Repeat the process swapping the roles of nums1 and nums2, and sum the results.
• Step 6: Return the total count.

This approach leverages hash maps for O(1) lookup and avoids unnecessary recomputation, making it much more efficient than brute-force.

Example Walkthrough

Suppose nums1 = [7,4] and nums2 = [5,2,8,9].

• For a = 7 in nums1: target = 49.
  • Check all pairs (b, c) in nums2 where b * c = 49.
  • Possible pairs: (b=7, c=7), but 7 is not in nums2 so no pairs.
• For a = 4: target = 16.
  • Possible pairs: (b=2, c=8) and (b=8, c=2).
  • Both 2 and 8 are in nums2. There is 1 occurrence of each, so 1 pair.
• Now swap roles: for each a in nums2, check pairs in nums1.
  • For a = 8: target = 64.
  • Possible pairs: (b=8, c=8) but 8 is not in nums1.
  • For a = 9: target = 81.
  • Possible pairs: (b=9, c=9), not in nums1.
  • For a = 2: target = 4.
  • Possible pairs: (b=2, c=2), not in nums1.
  • For a = 5: target = 25.
  • Possible pairs: (b=25, c=1), neither in nums1.
• Final answer: Only one valid triplet: (a=4 in nums1, b=2, c=8 in nums2).

Time and Space Complexity

• Brute-force: For each element in one array, try every pair in the other array. That is O(N * M^2) time, where N and M are the lengths of the two arrays.
• Optimized approach:
  • Building the frequency map: O(M) for one array.
  • For each of N elements, iterating over all unique values in the other array (let's say K unique values): O(N * K).
  • Since K ≤ M, total time is O(N * K + M) per direction, so O(N * K + M * L) overall (where L is the number of unique values in the other array).
  • Space complexity: O(M) for the hash map.
• In the worst case (all unique numbers), this is O(N * M), much better than O(N * M^2).

Summary

The key to solving this problem efficiently is recognizing that counting pairs with a product equal to a square can be reduced to counting frequencies of factors, which is ideal for hash maps. By iterating over one array and using the frequency map of the other, we avoid redundant work and achieve a much faster solution than brute-force. The approach is both simple and powerful, making good use of hash maps and careful counting logic to avoid overcounting.