BinarySearch 0979 Bounded Square Sums

Problem statement

https://binarysearch.com/problems/Bounded-Square-Sums/

Solution

Evaluate squares and sort, then use binary search or two pointers.

Complexity

It is O(n log n + m log m + m log n) for time and O(m + n) for space.

Code

class Solution:
    def solve(self, A, B, l, u):
        A = sorted(x*x for x in A)
        B = sorted(x*x for x in B)
        ans = 0
        for x in A:
            idx1 = bisect.bisect_left(B, l - x)
            idx2 = bisect.bisect(B, u - x)
            ans += (idx2 - idx1)
        return ans