Number of ways where square of number is equal to product of two numbers

Array Hash Table Math Two Pointers

43.2% acceptance

Feb 25, 2026

397

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.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn num_triplets(nums1: Vec<i32>, nums2: Vec<i32>) -> i32 {
    fn count(sq: &Vec<i32>, arr: &Vec<i32>) -> i32 {
      // Count triplets where sq[i]^2 == arr[j] * arr[k], j < k
      let mut total = 0;
      use std::collections::HashMap;
      for &s in sq {
        let target = (s as i64) * (s as i64);
        let mut freq: HashMap<i64, i64> = HashMap::new();
        for &a in arr {
          let a = a as i64;
          if target % a == 0 {
            let need = target / a;
            total += freq.get(&need).copied().unwrap_or(0);
          }
          *freq.entry(a).or_insert(0) += 1;
        }
      }
      total as i32
    }
    count(&nums1, &nums2) + count(&nums2, &nums1)
  }
}