Pythonで各クエリを含む最小区間のサイズを効率的に求めるプログラム

問題概要

区間のリスト intervals が与えられます。intervals[i] はペア (left_i, right_i) で表され、i 番目の区間が left_i で始まり right_i で終わることを意味します(両端を含みます)。同時に、queries という別の配列も与えられます。

j 番目のクエリに対する答えは、「left_i ≤ queries[j] ≤ right_i」を満たす区間の中で最もサイズ(長さ)が小さい区間のサイズです。条件を満たす区間が存在しない場合は -1 を返します。すべてのクエリへの回答を配列として求めてください。

具体例

たとえば、入力が intervals = [[2,5],[3,5],[4,7],[5,5]]、queries = [3,4,5,6] のとき、出力は [3, 3, 1, 4] になります。各クエリは次のように処理されます。

• query = 3 の場合:3 を含む最小の区間は [3,5] なので、5 − 3 + 1 = 3。
• query = 4 の場合:4 を含む最小の区間は [3,5] なので、5 − 3 + 1 = 3。
• query = 5 の場合:5 を含む最小の区間は [5,5] なので、5 − 5 + 1 = 1。
• query = 6 の場合:6 を含む最小の区間は [4,7] なので、7 − 4 + 1 = 4。

解き方の考え方

この問題は「ソート」と「優先度付きキュー(ヒープ)」を組み合わせると効率よく解けます。基本的なアイデアは次のとおりです。

1. クエリを昇順にソートし、小さい値から順に処理します。
2. 区間も開始位置で並べ替えておき、現在のクエリ値以降で始まる区間を順次ヒープに追加します。
3. ヒープには「(区間のサイズ, 区間の終了位置)」を格納し、サイズが小さいものが常に先頭に来るようにします。
4. すでに現在のクエリを含まない(終了位置がクエリより前の)区間はヒープから取り除きます。
5. こうすることで、ヒープの先頭にある要素が「現在のクエリを含む最小の区間」となります。

アルゴリズムの手順

• intervals を降順にソートします(末尾から開始位置の小さい区間を取り出せるようにするため)。
• h := 空のリスト(最小ヒープとして使用)
• res := 空の辞書(クエリごとの答えを保存)
• ソート済みの queries の各 q について、以下を繰り返します。
  • intervals が空でなく、末尾の区間の開始位置が q 以下である間:
    • (i, j) := intervals の末尾から区間を取り出して削除
    • j ≥ q(区間の終了位置が q 以上)なら、ペア (j − i + 1, j) をヒープ h に挿入
  • h が空でなく、ヒープ先頭の区間の終了位置が q 未満である間、先頭要素を削除(期限切れの区間を除去)
  • res[q] := h が空でなければヒープ先頭のサイズ、空なら -1
• queries の元の順序どおりに res[q] を並べたリストを返します。

Python実装例

以下の実装を見ると理解が深まります。

import heapq

def solve(intervals, queries):
    intervals = sorted(intervals)[::-1]
    h = []
    res = {}
    for q in sorted(queries):
        while intervals and intervals[-1][0] <= q:
            i, j = intervals.pop()
            if j >= q:
                heapq.heappush(h, [j - i + 1, j])
        while h and h[0][1] < q:
            heapq.heappop(h)
        res[q] = h[0][0] if h else -1
    return [res[q] for q in queries]

intervals = [[2,5],[3,5],[4,7],[5,5]]
queries = [3,4,5,6]
print(solve(intervals, queries))

入力

[[2,5],[3,5],[4,7],[5,5]], [3,4,5,6]

出力

[3, 3, 1, 4]

計算量

区間の数を n、クエリの数を m とすると、ソートに O(n log n + m log m) かかり、各区間・各クエリはそれぞれ高々1回ずつヒープに出入りするため、全体の計算量は O((n + m) log(n + m)) となります。各クエリごとに全区間を調べる素朴な O(n × m) の手法と比べて大幅に高速であり、大規模な入力にも対応できます。