Python程序：查找平方数等于两个数乘积的方式数量

假设我们有两个数组 nums1 和 nums2，我们需要找到形成的三元组（类型 1 和类型 2）的数量，遵循以下两个规则：

  • 三元组 (i, j, k) 如果 nums1[i]^2 = nums2[j] * nums2[k]，其中 [0 <= i < nums1 的大小 且 0 <= j < k < nums2 的大小]。
  • 三元组 (i, j, k) 如果 nums2[i]^2 = nums1[j] * nums1[k]，其中 [0 <= i < nums2 的大小 且 0 <= j < k < nums1 的大小]。

因此，如果输入类似于 nums1 = [7,4] nums2 = [5,2,8,9]，则输出将为 1，因为存在一个类型 1 的三元组，(1,1,2)，nums1[1]^2 = nums2[1] * nums2[2] = (16 = 2*8)。

为了解决这个问题，我们将遵循以下步骤

  • cnt1 := 一个映射，用于保存 nums1 中每个元素及其计数
  • cnt2 := 一个映射，用于保存 nums2 中每个元素及其计数
  • 定义一个函数 triplets()。它将接收 arr1 和 arr2 作为参数
  • ans := 0
  • 对于 arr1 的每个元素 t 及其值 v，执行以下操作：
    • k := 如果 arr2 中存在 t，则为 arr2[t]，否则为 0
    • tmp := k*(k - 1) / 2
    • sq := t * t
    • 对于 arr2 中的每个元素 m，执行以下操作：
      • 如果 m < t 且 sq 除以 m 的余数为 0，则
        • tmp := tmp + (如果 arr2 中存在 m，则为 arr2[m]，否则为 0) * (如果 arr2 中存在 sq/m 的商，则为 arr2[sq/m 的商]，否则为 0)
  • ans := ans + tmp * v
  • 返回 ans
  • 从主方法中，执行以下操作：
  • 返回 triplets(cnt1, cnt2) + triplets(cnt2, cnt1)

示例

让我们看看以下实现以更好地理解：

from collections import Counter
def solve(nums1, nums2):
   cnt1 = Counter(nums1)
   cnt2 = Counter(nums2)

   def triplets(arr1, arr2):
      ans = 0
      for t, v in arr1.items():
         k = arr2.get(t, 0)
         tmp = k * (k - 1) // 2
         sq = t * t
         for m in arr2:
            if m < t and sq % m == 0:
               tmp += arr2.get(m, 0) * arr2.get(sq // m, 0)
            ans += tmp * v
      return ans

   return triplets(cnt1, cnt2) + triplets(cnt2, cnt1)
nums1 = [7,4]
nums2 = [5,2,8,9]
print(solve(nums1, nums2))

输入

[7,4],[5,2,8,9]

输出

2

更新于: 2021年10月4日

163 次查看

开启你的 职业生涯

通过完成课程获得认证

开始学习
广告