Chương trình đếm cặp số bình phương trong khoảng cho trước Python

Cho hai mảng nums1, nums2 và hai số lower, upper. Đếm số cặp (i,j) thỏa lower ≤ nums1[i]² + nums2[j]² ≤ upper.

Ví dụ: nums1=[5,3,2], nums2=[8,12,6], lower=10, upper=50 → 2 (cặp (3,6) và (2,6)).

Giải thuật (Bình phương + Tìm kiếm nhị phân)

1. Bình phương tất cả phần tử hai mảng
2. Sắp xếp mảng lớn hơn (giả sử nums2)
3. Với mỗi e1 trong nums1: tìm số phần tử trong nums2 nằm trong [lower-e1, upper-e1] bằng bisect_left/right
4. Tổng cộng số lượng

Ví dụ triển khai

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.sort()
    res = 0
    for e1 in nums1:
        st = bisect_left(nums2, lower - e1)
        en = bisect_right(nums2, upper - e1)
        res += en - st
    return res

print(solve([5, 3, 2], [8, 12, 6], 10, 50))

Đầu vào

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

Đầu ra

2