Python程序：查找元素平方在给定范围内的对数

假设我们有两个数字列表nums1和nums2，以及两个数字lower和upper。我们需要找到这样的对数(i, j)，使得lower ≤ nums1[i]^2 + nums2[j]^2 ≤ upper。

例如，如果输入为nums1 = [5, 3, 2] nums2 = [8, 12, 6] lower = 10 upper = 50，则输出为2，因为对数为(1, 2)和(2, 2)。

• 10 <= 3^2 + 6^2 << 50 = 10 <= 45 << 50
• 10 <= 2^2 + 6^2 << 50 = 10 <= 40 << 50

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

• 将nums1中的每个元素替换为其平方。
• 将nums2中的每个元素替换为其平方。
• n := nums1的大小
• m := nums2的大小
• 如果n > m，则
  • 交换nums1和nums2
  • 交换n和m
• nums2 := 对列表nums2进行排序
• res := 0
• 对于nums1中的每个元素e1：
  • st := 在nums2中插入(lower - e1)的最左位置，以便元素保持排序。
  • en := 在nums2中插入(upper - e1)的最右位置，以便元素保持排序。
  • count := en - st
  • res := res + count
• 返回res

示例

让我们来看下面的实现，以便更好地理解：

from bisect import bisect_left, bisect_right
def solve(nums1, nums2, lower, upper):
   nums1 = [i * i for i in nums1]
   nums2 = [i * i for i in nums2]
   n, m = len(nums1), len(nums2)
   if n > m:
      nums1, nums2 = nums2, nums1
      n, m = m, n
   nums2 = sorted(nums2)
   res = 0
   for e1 in nums1:
      st = bisect_left(nums2, lower - e1)
      en = bisect_right(nums2, upper - e1)
      count = en - st
      res += count
   return res

nums1 = [5, 3, 2]
nums2 = [8, 12, 6]
lower = 10
upper = 50
print(solve(nums1, nums2, lower, upper))

输入

[5, 3, 2], [8, 12, 6], 10, 50

输出

2

Arnab Chakraborty

更新于：2021年10月16日

浏览量：184

开启你的职业生涯

完成课程获得认证

开始学习