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 次檢視

開啟你的 職業生涯

透過完成課程獲得認證

開始學習
廣告