Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 쿼리별 가장 가까운 방 찾기: 정렬과 이진 탐색 활용법

rooms라는 배열이 있다고 가정해 보겠습니다. rooms[i]는 [roomId_i, size_i] 쌍을 담고 있으며, 이는 id가 roomId_i이고 크기가 size_i인 방 하나를 나타냅니다. 모든 방 번호는 서로 중복되지 않습니다. 또한 queries라는 배열이 있는데, queries[j]는 [preferred_j, minSize_j] 쌍을 담고 있습니다.

j번째 쿼리에 대한 답은 다음 두 조건을 동시에 만족하는 방의 번호(id)입니다.

  • 방의 크기가 minSize_j 이상일 것
  • |id − preferred_j| 값이 최소화될 것

절대 차이가 동률이라면 더 작은 id를 가진 방을 선택합니다. 조건을 만족하는 방이 없다면 -1을 반환합니다. 즉, queries와 길이가 같은 answer 배열을 구성해야 하며, 각 원소에는 해당 쿼리의 답이 들어갑니다.

예시로 이해하기

입력이 rooms = [[2,2],[1,2],[3,2]], queries = [[3,1],[3,3],[5,2]]라고 한다면 출력은 [3, -1, 3]입니다. 그 이유는 다음과 같습니다.

  • 쿼리 [3,1]: |3 − 3| = 0이므로 3번 방이 가장 가깝고, 크기 2는 최소 크기 1 이상이므로 답은 3입니다.
  • 쿼리 [3,3]: 크기가 3 이상인 방이 존재하지 않으므로 답은 -1입니다.
  • 쿼리 [5,2]: |3 − 5| = 2이므로 3번 방이 가장 가깝고, 크기 2는 최소 크기 2 이상이므로 답은 3입니다.

해결 접근 방식

이 문제를 효율적으로 해결하려면 다음 단계를 따릅니다.

  • rooms를 크기 기준으로 오름차순 정렬하고, 크기가 같으면 방 id를 기준으로 정렬합니다.
  • queries를 (qid, size, i) 형태의 튜플 리스트로 변환합니다. 여기서 i는 원래 쿼리의 인덱스입니다.
  • 쿼리를 크기(size) 기준 내림차순으로 정렬하고, 크기가 같으면 preferred 기준, 둘 다 같으면 인덱스 기준으로 정렬합니다.
  • ans := queries와 같은 길이의 배열을 만들고 -1로 초기화합니다.
  • X := 새로운 빈 리스트를 준비합니다.

이제 각 쿼리 (qid, size, i)에 대해 다음을 수행합니다.

  • rooms가 비어 있지 않고, 마지막 방의 크기가 현재 쿼리의 size보다 크거나 같은 동안 반복합니다.
    • (idr, p) := rooms에서 마지막 요소를 꺼냅니다.
    • idr을 X에 삽입하면서 정렬 상태를 유지합니다(bisect.insort 활용).
  • X가 비어 있지 않다면:
    • j := qid를 삽입했을 때 X의 정렬 상태가 유지되는 위치(bisect.bisect 결과)
    • j == len(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]

모든 쿼리를 처리한 후 ans를 반환합니다. 이 방식은 방과 쿼리를 한 번씩만 순회하며 이진 탐색으로 최적 후보를 찾기 때문에 전체 시간 복잡도가 O((N + Q) log(N + Q)) 수준으로 매우 효율적입니다.

구현 예제

아래 코드를 통해 실제 동작을 확인해 보겠습니다.

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]