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

Description

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]. (42 = 2 * 8). 

Example 2:

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:

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] <= 105

Solutions

Solution 1: Hash Table + Enumeration

We use a hash table \(\textit{cnt1}\) to count the occurrences of each pair \((\textit{nums}[j], \textit{nums}[k])\) in \(\textit{nums1}\), where \(0 \leq j < k < m\), and \(m\) is the length of the array \(\textit{nums1}\). Similarly, we use a hash table \(\textit{cnt2}\) to count the occurrences of each pair \((\textit{nums}[j], \textit{nums}[k])\) in \(\textit{nums2}\), where \(0 \leq j < k < n\), and \(n\) is the length of the array \(\textit{nums2}\).

Next, we enumerate each number \(x\) in the array \(\textit{nums1}\) and calculate the value of \(\textit{cnt2}[x^2]\), which is the number of pairs \((\textit{nums}[j], \textit{nums}[k])\) in \(\textit{nums2}\) that satisfy \(\textit{nums}[j] \times \textit{nums}[k] = x^2\). Similarly, we enumerate each number \(x\) in the array \(\textit{nums2}\) and calculate the value of \(\textit{cnt1}[x^2]\), which is the number of pairs \((\textit{nums}[j], \textit{nums}[k])\) in \(\textit{nums1}\) that satisfy \(\textit{nums}[j] \times \textit{nums}[k] = x^2\). Finally, we return the sum of the two results.

The time complexity is \(O(m^2 + n^2 + m + n)\), and the space complexity is \(O(m^2 + n^2)\). Here, \(m\) and \(n\) are the lengths of the arrays \(\textit{nums1}\) and \(\textit{nums2}\), respectively.

Python3JavaC++GoTypeScript

class Solution:
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        def count(nums: List[int]) -> Counter:
            cnt = Counter()
            for j in range(len(nums)):
                for k in range(j + 1, len(nums)):
                    cnt[nums[j] * nums[k]] += 1
            return cnt

        def cal(nums: List[int], cnt: Counter) -> int:
            return sum(cnt[x * x] for x in nums)

        cnt1 = count(nums1)
        cnt2 = count(nums2)
        return cal(nums1, cnt2) + cal(nums2, cnt1)

class Solution {
    public int numTriplets(int[] nums1, int[] nums2) {
        var cnt1 = count(nums1);
        var cnt2 = count(nums2);
        return cal(cnt1, nums2) + cal(cnt2, nums1);
    }

    private Map<Long, Integer> count(int[] nums) {
        Map<Long, Integer> cnt = new HashMap<>();
        int n = nums.length;
        for (int j = 0; j < n; ++j) {
            for (int k = j + 1; k < n; ++k) {
                long x = (long) nums[j] * nums[k];
                cnt.merge(x, 1, Integer::sum);
            }
        }
        return cnt;
    }

    private int cal(Map<Long, Integer> cnt, int[] nums) {
        int ans = 0;
        for (int x : nums) {
            long y = (long) x * x;
            ans += cnt.getOrDefault(y, 0);
        }
        return ans;
    }
}

class Solution {
public:
    int numTriplets(vector<int>& nums1, vector<int>& nums2) {
        auto cnt1 = count(nums1);
        auto cnt2 = count(nums2);
        return cal(cnt1, nums2) + cal(cnt2, nums1);
    }

    unordered_map<long long, int> count(vector<int>& nums) {
        unordered_map<long long, int> cnt;
        for (int i = 0; i < nums.size(); i++) {
            for (int j = i + 1; j < nums.size(); j++) {
                cnt[(long long) nums[i] * nums[j]]++;
            }
        }
        return cnt;
    }

    int cal(unordered_map<long long, int>& cnt, vector<int>& nums) {
        int ans = 0;
        for (int x : nums) {
            ans += cnt[(long long) x * x];
        }
        return ans;
    }
};

func numTriplets(nums1 []int, nums2 []int) int {
    cnt1 := count(nums1)
    cnt2 := count(nums2)
    return cal(cnt1, nums2) + cal(cnt2, nums1)
}

func count(nums []int) map[int]int {
    cnt := map[int]int{}
    for j, x := range nums {
        for _, y := range nums[j+1:] {
            cnt[x*y]++
        }
    }
    return cnt
}

func cal(cnt map[int]int, nums []int) (ans int) {
    for _, x := range nums {
        ans += cnt[x*x]
    }
    return
}

function numTriplets(nums1: number[], nums2: number[]): number {
    const cnt1 = count(nums1);
    const cnt2 = count(nums2);
    return cal(cnt1, nums2) + cal(cnt2, nums1);
}

function count(nums: number[]): Map<number, number> {
    const cnt: Map<number, number> = new Map();
    for (let j = 0; j < nums.length; ++j) {
        for (let k = j + 1; k < nums.length; ++k) {
            const x = nums[j] * nums[k];
            cnt.set(x, (cnt.get(x) || 0) + 1);
        }
    }
    return cnt;
}

function cal(cnt: Map<number, number>, nums: number[]): number {
    return nums.reduce((acc, x) => acc + (cnt.get(x * x) || 0), 0);
}

Solution 2: Hash Table + Enumeration Optimization

We use a hash table \(\textit{cnt1}\) to count the occurrences of each number in \(\textit{nums1}\), and a hash table \(\textit{cnt2}\) to count the occurrences of each number in \(\textit{nums2}\).

Next, we enumerate each number \(x\) in the array \(\textit{nums1}\), and then enumerate each pair \((y, v1)\) in \(\textit{cnt2}\), where \(y\) is the key of \(\textit{cnt2}\) and \(v1\) is the value of \(\textit{cnt2}\). We calculate \(z = x^2 / y\). If \(y \times z = x^2\), and if \(y = z\), it means \(y\) and \(z\) are the same number, then the number of ways to choose two numbers from \(v1\) is \(v1 \times (v1 - 1) = v1 \times (v2 - 1)\). If \(y \neq z\), then the number of ways to choose two numbers from \(v1\) is \(v1 \times v2\). Finally, we sum all the ways and divide by \(2\). The division by \(2\) is because we count the number of ways for the pair \((j, k)\), but \((j, k)\) and \((k, j)\) are the same way.

The time complexity is \(O(m \times n)\), and the space complexity is \(O(m + n)\). Here, \(m\) and \(n\) are the lengths of the arrays \(\textit{nums1}\) and \(\textit{nums2}\), respectively.

Python3JavaGoTypeScript

class Solution:
    def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
        def cal(nums: List[int], cnt: Counter) -> int:
            ans = 0
            for x in nums:
                for y, v1 in cnt.items():
                    z = x * x // y
                    if y * z == x * x:
                        v2 = cnt[z]
                        ans += v1 * (v2 - int(y == z))
            return ans // 2

        cnt1 = Counter(nums1)
        cnt2 = Counter(nums2)
        return cal(nums1, cnt2) + cal(nums2, cnt1)

class Solution {
    public int numTriplets(int[] nums1, int[] nums2) {
        var cnt1 = count(nums1);
        var cnt2 = count(nums2);
        return cal(cnt1, nums2) + cal(cnt2, nums1);
    }

    private Map<Integer, Integer> count(int[] nums) {
        Map<Integer, Integer> cnt = new HashMap<>();
        for (int x : nums) {
            cnt.merge(x, 1, Integer::sum);
        }
        return cnt;
    }

    private int cal(Map<Integer, Integer> cnt, int[] nums) {
        long ans = 0;
        for (int x : nums) {
            for (var e : cnt.entrySet()) {
                int y = e.getKey(), v1 = e.getValue();
                int z = (int) (1L * x * x / y);
                if (y * z == x * x) {
                    int v2 = cnt.getOrDefault(z, 0);
                    ans += v1 * (y == z ? v2 - 1 : v2);
                }
            }
        }
        return (int) (ans / 2);
    }
}

func numTriplets(nums1 []int, nums2 []int) int {
    cnt1 := count(nums1)
    cnt2 := count(nums2)
    return cal(cnt1, nums2) + cal(cnt2, nums1)
}

func count(nums []int) map[int]int {
    cnt := map[int]int{}
    for _, x := range nums {
        cnt[x]++
    }
    return cnt
}

func cal(cnt map[int]int, nums []int) (ans int) {
    for _, x := range nums {
        for y, v1 := range cnt {
            z := x * x / y
            if y*z == x*x {
                if v2, ok := cnt[z]; ok {
                    if y == z {
                        v2--
                    }
                    ans += v1 * v2
                }
            }
        }
    }
    ans /= 2
    return
}

function numTriplets(nums1: number[], nums2: number[]): number {
    const cnt1 = count(nums1);
    const cnt2 = count(nums2);
    return cal(cnt1, nums2) + cal(cnt2, nums1);
}

function count(nums: number[]): Map<number, number> {
    const cnt: Map<number, number> = new Map();
    for (const x of nums) {
        cnt.set(x, (cnt.get(x) || 0) + 1);
    }
    return cnt;
}

function cal(cnt: Map<number, number>, nums: number[]): number {
    let ans: number = 0;
    for (const x of nums) {
        for (const [y, v1] of cnt) {
            const z = Math.floor((x * x) / y);
            if (y * z == x * x) {
                const v2 = cnt.get(z) || 0;
                ans += v1 * (y === z ? v2 - 1 : v2);
            }
        }
    }
    return ans / 2;
}

Comments