Number of Ways Where Square of Number Is Equal to Product of Two Numbers

• The problem can be broken down into two symmetric subproblems: counting Type 1 triplets and counting Type 2 triplets. The total count is the sum of these two.
• We can create a helper function, say count(arr1, arr2), which counts triplets where a square of a number from arr1 equals the product of two numbers from arr2. The final answer would be count(nums1, nums2) + count(nums2, nums1).
• The count function iterates through every element arr1[i]. For each arr1[i], it then iterates through all possible pairs (arr2[j], arr2[k]) with j < k.
• Inside the innermost loop, it checks if (long)arr1[i] * arr1[i] == (long)arr2[j] * arr2[k]. If the condition is true, a counter is incremented.
• Using long for calculations is crucial to prevent integer overflow, as numbers can be up to 10^5, and their square can be up to 10^10.

The core idea is to check every single combination of indices (i, j, k) that satisfy the problem's constraints. We can implement a helper function that takes two arrays, arr1 and arr2, and counts the number of triplets where the square of an element from arr1 equals the product of two distinct elements from arr2. By calling this helper function twice with the arguments swapped (count(nums1, nums2) and count(nums2, nums1)), we can find the total number of valid triplets.

1class Solution {2    public int numTriplets(int[] nums1, int[] nums2) {3        return countTriplets(nums1, nums2) + countTriplets(nums2, nums1);4    }5 6    private int countTriplets(int[] arr1, int[] arr2) {7        int count = 0;8        int n = arr1.length;9        int m = arr2.length;10        if (m < 2) {11            return 0;12        }13        for (int i = 0; i < n; i++) {14            long square = (long) arr1[i] * arr1[i];15            for (int j = 0; j < m; j++) {16                for (int k = j + 1; k < m; k++) {17                    if (square == (long) arr2[j] * arr2[k]) {18                        count++;19                    }20                }21            }22        }23        return count;24    }25}

• Define a helper function count(arr1, arr2).
• Sort arr2.
• Initialize total_triplets = 0.
• For each num1 in arr1:
  • Calculate target = (long)num1 * num1.
  • Initialize two pointers on arr2: left = 0, right = arr2.length - 1.
  • While left < right:
    • Calculate product = (long)arr2[left] * arr2[right].
    • If product < target, increment left.
    • Else if product > target, decrement right.
    • Else (product == target):
      • If arr2[left] == arr2[right], all elements between them are the same. The number of pairs is k * (k-1) / 2 where k is right - left + 1. Add this to the count and break the inner loop.
      • If arr2[left] != arr2[right], count the number of duplicates for arr2[left] and arr2[right], multiply these counts, add to the total, and move the pointers past the duplicates.
• Return total_triplets.
• The main function returns count(nums1, nums2) + count(nums2, nums1).

We can reframe the problem for a fixed num1 from the first array as a 'Two Sum'-like problem on the second array. We need to find pairs (num2_j, num2_k) such that their product is num1^2. If the second array is sorted, this can be solved efficiently using two pointers.

The overall strategy is to have a helper function that first sorts the second array (arr2). Then, for each element in the first array (arr1), it uses a two-pointer scan (left and right) on arr2 to find pairs whose product matches the target square. A key part of this approach is correctly handling duplicate numbers to ensure every valid pair is counted exactly once.

1import java.util.Arrays;2 3class Solution {4    public int numTriplets(int[] nums1, int[] nums2) {5        return countTriplets(nums1, nums2) + countTriplets(nums2, nums1);6    }7 8    private int countTriplets(int[] arr1, int[] arr2) {9        Arrays.sort(arr2);10        int count = 0;11        int n = arr1.length;12        int m = arr2.length;13        if (m < 2) {14            return 0;15        }16 17        for (int i = 0; i < n; i++) {18            long target = (long) arr1[i] * arr1[i];19            int left = 0, right = m - 1;20            while (left < right) {21                long product = (long) arr2[left] * arr2[right];22                if (product < target) {23                    left++;24                } else if (product > target) {25                    right--;26                } else { // product == target27                    if (arr2[left] == arr2[right]) {28                        int numElements = right - left + 1;29                        count += numElements * (numElements - 1) / 2;30                        break; // All elements in between are the same31                    }32                    int leftCount = 1;33                    while (left + 1 < right && arr2[left] == arr2[left + 1]) {34                        leftCount++;35                        left++;36                    }37                    int rightCount = 1;38                    while (right - 1 > left && arr2[right] == arr2[right - 1]) {39                        rightCount++;40                        right--;41                    }42                    count += leftCount * rightCount;43                    left++;44                    right--;45                }46            }47        }48        return count;49    }50}

• Define a helper function count(arr1, arr2).
• Create a frequency map freq2 for all numbers in arr2.
• Initialize total_triplets = 0.
• For each num1 in arr1:
  • Calculate target = (long)num1 * num1.
  • For each unique number num2 and its count2 in freq2:
    • If target is divisible by num2:
      • Calculate required = target / num2.
      • If required exists in freq2:
        • If num2 == required, add count2 * (count2 - 1) / 2 to total_triplets.
        • If num2 < required, add count2 * freq2.get(required) to total_triplets to avoid double counting.
• Return total_triplets.
• The main function returns count(nums1, nums2) + count(nums2, nums1).

The key optimization here is to avoid re-computing information. For a given target square, we need to find pairs in the second array. Instead of a linear scan or a two-pointer scan for every target, we can count the occurrences of each number in the second array and store them in a hash map.

Then, for each number num1 in the first array, we iterate through the unique numbers num2 in the hash map. For each num2, we can calculate the required number that would satisfy num2 * required = num1^2. We can then check the hash map for the existence and frequency of this required number. Special care is taken to handle cases where num2 is equal to required (we need to choose two from the same group of numbers) and to avoid double-counting pairs like (a, b) and (b, a).

1import java.util.HashMap;2import java.util.Map;3 4class Solution {5    public int numTriplets(int[] nums1, int[] nums2) {6        return (int) (countTriplets(nums1, nums2) + countTriplets(nums2, nums1));7    }8 9    private long countTriplets(int[] arr1, int[] arr2) {10        Map<Integer, Integer> freq2 = new HashMap<>();11        for (int num : arr2) {12            freq2.put(num, freq2.getOrDefault(num, 0) + 1);13        }14 15        long count = 0;16        for (int num1 : arr1) {17            long target = (long) num1 * num1;18            for (Map.Entry<Integer, Integer> entry : freq2.entrySet()) {19                long num2 = entry.getKey();20                long count2 = entry.getValue();21 22                if (target % num2 != 0) {23                    continue;24                }25                long requiredNum = target / num2;26                27                if (freq2.containsKey((int)requiredNum)) {28                    if (num2 == requiredNum) {29                        count += count2 * (count2 - 1) / 2;30                    } else if (num2 < requiredNum) {31                        count += count2 * freq2.get((int)requiredNum);32                    }33                }34            }35        }36        return count;37    }38}