Find Number of Pairs Where Elements Square is Within Given Range in Python

Suppose we have two list of numbers nums1 and nums2. And also have two numbers lower and upper. We have to find the number of pairs (i, j) such that lower ≤ nums1[i]^2 + nums2[j]^2 ≤ upper.

So, if the input is like nums1 = [5, 3, 2] nums2 = [8, 12, 6] lower = 10 upper = 50, then the output will be 2 because the pairs are like (1, 2) and (2, 2)

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

To solve this, we will follow these steps −

• replace each element by its square in nums1
• replace each element by its square in nums2
• n := size of nums1
• m := size of nums2
• if n > m , then
  • swap nums1 and nums2
  • swap n and m
• nums2 := sort the list nums2
• res := 0
• for each e1 in nums1, do
  • st := left most position to insert (lower - e1) into nums2 so that elements are sorted
  • en := right most position to insert (upper - e1) into nums2 so that elements are sorted
  • count := en - st
  • res := res + count
• return res

Example

Let us see the following implementation to get better understanding −

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))

Input

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

Output

2