Pythonで2つのリストの要素の2乗和が指定範囲内となるペアの数を効率的に求めるプログラム

問題概要

2つの数値リスト nums1 と nums2、および下限値 lower と上限値 upper が与えられます。このとき、次の条件を満たすペア (i, j) の総数を求めるのが目的です。

lower ≤ nums1[i]² + nums2[j]² ≤ upper

入力例

たとえば、nums1 = [5, 3, 2]、nums2 = [8, 12, 6]、lower = 10、upper = 50 という入力の場合、答えは 2 になります。条件を満たすのは (1, 2) と (2, 2) の2組だからです。

• 10 ≤ 3² + 6² = 45 ≤ 50
• 10 ≤ 2² + 6² = 40 ≤ 50

解き方のポイント

すべてのペア (i, j) を総当たりで調べると計算量が O(n × m) となり、リストが大きい場合には非効率です。そこで、二分探索(bisect) を活用して高速化します。手順は以下の通りです。

1. nums1、nums2 の各要素をそれぞれ自乗した値に置き換える。
2. n を nums1 のサイズ、m を nums2 のサイズとする。
3. n > m の場合は、nums1 と nums2(および n と m)を入れ替える。短い方を走査対象にすることで処理を効率化するためです。
4. nums2 を昇順にソートする。
5. 結果を格納する変数 res を 0 で初期化する。
6. nums1 の各要素 e1 について次を繰り返す。
   • st:nums2 に (lower − e1) を挿入してもソート順が保たれる最左位置(bisect_left)
   • en:nums2 に (upper − e1) を挿入してもソート順が保たれる最右位置(bisect_right)
   • count = en − st を計算し、res に加算する。
7. 最後に res を返す。

bisect_left は等しい要素の「手前」の位置を返し、bisect_right は「後ろ」の位置を返します。この2つの差分を取ることで、lower 以上 upper 以下の範囲に収まる要素の個数を、重複値を含めて正確に数えられるのがこの手法の肝です。

Pythonでの実装例

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

入力

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

出力

2

計算量について

nums2 のソートに O(m log m)、各要素ごとの二分探索に O(log m) かかるため、全体の計算量は O((n + m) log m) となります。全ペアを直接調べる O(n × m) の素朴なアプローチと比べて大幅に高速であり、リストのサイズが大きいケースでも実用的な性能を発揮します。