Skip to content

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

  • Time: $O(|\texttt{nums1}||\texttt{nums2}|)$
  • Space: $O(|\texttt{nums1}| + |\texttt{nums2}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution {
 public:
  int numTriplets(vector<int>& nums1, vector<int>& nums2) {
    return countTriplets(nums1, nums2) + countTriplets(nums2, nums1);
  }

 private:
  // Returns the number of triplet (i, j, k) if A[i]^2 == B[j] * B[k].
  int countTriplets(const vector<int>& A, const vector<int>& B) {
    int res = 0;
    unordered_map<int, int> count;

    for (const int b : B)
      ++count[b];

    for (const int a : A) {
      const long target = static_cast<long>(a) * a;
      for (const auto& [b, freq] : count) {
        if (target % b > 0 || !count.count(target / b))
          continue;
        if (target / b == b)
          res += freq * (freq - 1);
        else
          res += freq * count[target / b];
      }
    }

    return res / 2;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution {
  public int numTriplets(int[] nums1, int[] nums2) {
    return countTriplets(nums1, nums2) + countTriplets(nums2, nums1);
  }

  // Returns the number of triplet (i, j, k) if A[i]^2 == B[j] * B[k].
  private int countTriplets(int[] A, int[] B) {
    int res = 0;
    Map<Integer, Integer> count = new HashMap<>();

    for (final int b : B)
      count.merge(b, 1, Integer::sum);

    for (final int a : A) {
      long target = (long) a * a;
      for (Map.Entry<Integer, Integer> entry : count.entrySet()) {
        final int b = entry.getKey();
        final int freq = entry.getValue();
        if (target % b > 0 || !count.containsKey((int) (target / b)))
          continue;
        if (target / b == b)
          res += freq * (freq - 1);
        else
          res += freq * count.get((int) (target / b));
      }
    }

    return res / 2;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
  def numTriplets(self, nums1: List[int], nums2: List[int]) -> int:
    def countTriplets(A: List[int], B: List[int]):
      """Returns the number of triplet (i, j, k) if A[i]^2 == B[j] * B[k]."""
      res = 0
      count = collections.Counter(B)

      for a in A:
        target = a * a
        for b, freq in count.items():
          if target % b > 0 or target // b not in count:
            continue
          if target // b == b:
            res += freq * (freq - 1)
          else:
            res += freq * count[target // b]

      return res // 2

    return countTriplets(nums1, nums2) + countTriplets(nums2, nums1)