Pythonでクエリから最も近い部屋を見つける効率的なアルゴリズム

問題概要

roomsという配列があるとします。rooms[i]はペア[roomId_i, size_i]を含み、roomId_iは部屋のID、size_iはその部屋のサイズを表します。すべての部屋番号は一意(重複なし)です。

さらに、queriesという別の配列も与えられます。queries[j]はペア[preferred_j, minSize_j]を含みます。j番目のクエリに対する答えは、次の条件を満たす部屋のIDです。

• 部屋のサイズがminSize_j以上であること

• |id − preferred_j| が最小となること

絶対差が同点になった場合は、IDが最も小さい部屋を採用します。条件を満たす部屋が存在しない場合は−1を返します。このとき、queriesと同じ長さのanswer配列を作成し、j番目の要素にj番目のクエリへの回答を格納して返すのが目的です。

具体例で理解する

たとえば、rooms = [[2,2],[1,2],[3,2]]、queries = [[3,1],[3,3],[5,2]] が入力された場合、出力は [3, -1, 3] になります。理由は次のとおりです。

• クエリ [3,1] の場合: 部屋3が最も近く(|3 − 3| = 0)、そのサイズ2はminSizeの1以上であるため、答えは3になります。

• クエリ [3,3] の場合: サイズが3以上の部屋は存在しないため、答えは−1になります。

• クエリ [5,2] の場合: 部屋3が最も近く(|3 − 5| = 2)、そのサイズ2はminSizeの2以上であるため、答えは3になります。

解法のアプローチ

この問題を効率的に解くには、次の手順に従います。

1. roomsをサイズの昇順でソートします。サイズが同じ場合は部屋IDの昇順で並べます。

2. queriesを、インデックスiを含むタプル(qid, size, i)のリストに変換します。

3. queriesをサイズの降順でソートします。サイズが同じ場合はpreferredの降順、それも同じ場合はインデックス順に並べます。

4. ans := queriesと同じ長さの配列を用意し、すべて−1で初期化します。

5. X := 新しい空のリストを用意します。

6. queries内の各(qid, size, i)に対して、以下を実行します。

   • roomsが空でなく、roomsの末尾要素のサイズがsize以上である間、次を繰り返します。

     • (idr, p) := roomsの末尾から要素を取り除きます。

     • idrをXに挿入し、Xがソートされた状態を保ちます(bisect.insortを使用)。

   • Xが空でない場合:

     • j := qidを挿入してもXがソートされたままになる挿入位置(bisect.bisectで取得)

     • jがXの長さと等しい場合 → ans[i] := Xの末尾要素

     • jが0の場合 → ans[i] := X[0]

     • それ以外の場合:

       • X[j] − qid < qid − X[j−1] ならば ans[i] := X[j]

       • そうでなければ ans[i] := X[j−1]

7. ansを返します。

この手法のポイントは、クエリをminSizeの降順に処理することで、これまでに処理した部屋(サイズが現在のminSize以上のもの)だけをXに蓄積できる点にあります。こうすることで、各クエリの時点でXには「条件を満たす部屋のID」のみが含まれており、あとは二分探索で最も近いIDを高速に特定できます。

Pythonでの実装例

理解を深めるために、以下の実装例をご覧ください。

import bisect
def solve(rooms, queries):
   rooms.sort(key = lambda x: (x[1], x[0]))
   queries = [(qid,size,i) for i, (qid, size) in enumerate(queries)]
   queries.sort(key = lambda x: (x[1], x[0], x[2]), reverse = True)
   ans = [-1] * len(queries)
   X = []
   for qid, size, i in queries:
      while rooms and rooms[-1][1] >= size:
         idr, _ = rooms.pop()
         bisect.insort(X, idr)
      if X:
         j = bisect.bisect(X, qid)
         if j == len(X):
            ans[i] = X[-1]
         elif j == 0:
            ans[i] = X[0]
         else:
            if X[j] - qid < qid - X[j-1]:
               ans[i] = X[j]
            else:
               ans[i] = X[j-1]
   return ans

rooms = [[2,2],[1,2],[3,2]]
queries = [[3,1],[3,3],[5,2]]
print(solve(rooms, queries))

入力

[[2,2],[1,2],[3,2]], [[3,1],[3,3],[5,2]]

出力

[3, -1, 3]

計算量について

部屋とクエリの数をそれぞれn、mとすると、ソートにかかる時間はO(n log n + m log m)です。各部屋は一度だけXに挿入され、挿入・検索はいずれもO(log n)で行えるため、全体の計算量はO((n + m) log(n + m))となります。これは全件を毎回走査する素朴な方法(O(n × m))に比べて大幅に高速であり、大規模な入力にも対応できます。