Program to find number of pairs where elements square is within the given range in Python

Given two lists of numbers nums1 and nums2, and two boundary values lower and upper, we need to find the number of pairs (i, j) such that lower ? nums1[i]² + nums2[j]² ? upper.

Problem Understanding

For example, if nums1 = [5, 3, 2], nums2 = [8, 12, 6], lower = 10, and upper = 50, we need to check all possible combinations:

• Pair (1, 2): 3² + 6² = 9 + 36 = 45 (10 ? 45 ? 50) ?
• Pair (2, 2): 2² + 6² = 4 + 36 = 40 (10 ? 40 ? 50) ?

The answer is 2 valid pairs.

Algorithm Approach

The efficient solution uses binary search optimization:

• Square all elements in both lists
• Sort the longer list to enable binary search
• For each element in the shorter list, use binary search to count valid pairs
• Use bisect_left() and bisect_right() to find the range of valid elements

Implementation

from bisect import bisect_left, bisect_right

def count_pairs_in_range(nums1, nums2, lower, upper):
    # Square all elements
    nums1 = [i * i for i in nums1]
    nums2 = [i * i for i in nums2]
    
    n, m = len(nums1), len(nums2)
    
    # Optimize by making nums1 the smaller array
    if n > m:
        nums1, nums2 = nums2, nums1
        n, m = m, n
    
    # Sort the larger array for binary search
    nums2 = sorted(nums2)
    
    result = 0
    
    # For each element in smaller array
    for e1 in nums1:
        # Find range of valid elements in nums2
        # We need: lower <= e1 + e2 <= upper
        # So: (lower - e1) <= e2 <= (upper - e1)
        start = bisect_left(nums2, lower - e1)
        end = bisect_right(nums2, upper - e1)
        
        count = end - start
        result += count
    
    return result

# Test the function
nums1 = [5, 3, 2]
nums2 = [8, 12, 6]
lower = 10
upper = 50

print(f"Number of valid pairs: {count_pairs_in_range(nums1, nums2, lower, upper)}")

Number of valid pairs: 2

How It Works

The algorithm works by transforming the problem into a range search:

1. Preprocessing: Square all elements to avoid repeated calculations
2. Optimization: Always iterate through the smaller array to minimize operations
3. Binary Search: For each element e1 in the smaller array, find how many elements e2 in the sorted larger array satisfy lower - e1 ? e2 ? upper - e1
4. Counting: Use bisect_left() and bisect_right() to find the exact range of valid elements

Time Complexity

The time complexity is O(n log m + m log m) where n and m are the lengths of the input arrays. The space complexity is O(n + m) for storing the squared values.

Conclusion

This solution efficiently finds pairs whose sum of squares falls within a given range using binary search optimization. The key insight is sorting one array and using binary search to count valid pairs for each element in the other array.