Program to find closest room from queries in Python

Suppose there is an array called rooms. where rooms[i] contains a pair [roomId_i, size_i] denotes a room whose id is roomId_i and size is size_i. All room numbers are distinct. We also have another array queries, where queries[j] contains a pair [preferred_j, minSize_j]. The answer to the jth query is the room number id of a room such that −

  • The room has size of at least minSize_j, and

  • |id - preferred_j| is minimized.

Now, if there is a tie in the absolute difference, then use the room with the smallest id. If there is no such room, return -1. So we have to find an array called answer whose length is same as queries, that contains the answer to the jth query.

So, if the input is like rooms = [[2,2],[1,2],[3,2]] queries = [[3,1],[3,3],[5,2]], then the output will be [3, -1, 3] because

  • For query [3,1]: The room 3 is the closest because |3 - 3| = 0, and its size of 2 is at least 1, so the answer is 3.

  • For query [3,3]: There are no rooms whose size is at least 3, so the answer is -1.

  • For query [5,2]: The room 3 is the closest because |3 - 5| = 2, and its size of 2 is at least 2, so the answer is 3.

To solve this, we will follow these steps −

  • sort rooms based on size, when size are same then based on room id

  • queries = a list of pairs (qid,size,i) for index i, and pair (qid, size) in queries

  • sort the queries in reverse order based on size, if sizes are same, then based on preferred, when both are same then based on index

  • ans := an array of size same as size of queries and fill with -1

  • X := a new list

  • for each (qid, size, i) in queries, do

    • while rooms and size of last item of rooms >= size, do

      • (idr, p) := deleted last element from rooms

      • sort X after inserting idr

    • if X is not empty, then

      • j := index where to insert qid to remain X sorted

      • if j is same as size of X , then

        • ans[i] := last element of X

      • otherwise when j is same as 0, then

        • ans[i] := X[0]

      • otherwise,

        • if X[j] - qid < qid - X[j-1], then

          • ans[i] := X[j]

        • otherwise,

          • ans[i] := X[j-1]

  • return ans


Let us see the following implementation to get better understanding

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]
            if X[j] - qid < qid - X[j-1]:
               ans[i] = X[j]
               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]

Updated on: 08-Oct-2021


Kickstart Your Career

Get certified by completing the course

Get Started