Bounded Square Sums

Description

You are given two lists of integers a and b as well as integers lower and upper. Return the number of pairs (i, j) such that lower ≤ a[i] * a[i] + b[j] * b[j] ≤ upper.

Constraints:

• 0 ≤ n ≤ 100,000 where n is the length of a
• 0 ≤ m ≤ 100,000 where m is the length of b

Example 1

Input

a = [3, -1, 9]
b = [100, 5, -2]
lower = 7
upper = 99

Output

4

Explanation

We have the following pairs:

3 * 3 + 5 * 5 = 34
3 * 3 + -2 * -2 = 13
-1 * -1 + 5 * 5 = 26
9 * 9 + -2 * -2 = 85

---

Intuition

lower ≤ x ≤ upper → x ≤ upper - x≤ lower - 1

Algorithm Templates Binary Search, find target left

Implementation

import java.util.*;

class Solution {
    public int solve(int[] a, int[] b, int lower, int upper) {
        for (int i = 0; i < a.length; i++) {
            a[i] *= a[i];
        }
        for (int i = 0; i < b.length; i++) {
            b[i] *= b[i];
        }

        Arrays.sort(a);
        Arrays.sort(b);

        return helper(a, b, upper) - helper(a, b, lower - 1);
    }

    private int helper(int[] a, int[] b, int target) {
        int count = 0;
        for (int i = 0; i < a.length; i++) {
            int l = -1, r = b.length - 1;
            while (l < r) {
                int mid = l + r + 1 >> 1;
                if (a[i] + b[mid] <= target) {
                    l = mid;
                } else {
                    r = mid - 1;
                }
            }
            count += l + 1;
        }

        return count;
    }
}

Time Complexity

• Time: O(nlogn)
• Space: O(1)