Program to find closest room from queries in Python

Given an array of rooms where rooms[i] = [roomId_i, size_i] and an array of queries where queries[j] = [preferred_j, minSize_j], we need to find the closest room for each query based on specific criteria.

For each query, we need to find a room that:

• Has size of at least minSize_j

• |roomId - preferred_j| is minimized

• In case of tie, choose the room with smallest id

• Return -1 if no such room exists

Example

If rooms = [[2,2],[1,2],[3,2]] and queries = [[3,1],[3,3],[5,2]], the output is [3, -1, 3]:

• Query [3,1]: Room 3 is closest with |3-3| = 0 and size 2 ? 1

• Query [3,3]: No room has size ? 3, so return -1

• Query [5,2]: Room 3 is closest with |3-5| = 2 and size 2 ? 2

Algorithm

The approach uses sorting and binary search for efficient room matching:

1. Sort rooms by size (ascending), then by room id

2. Sort queries by size (descending) to process larger size requirements first

3. Maintain a sorted list of available room ids that meet current size requirement

4. Use binary search to find the closest room id for each query

Implementation

import bisect

def solve(rooms, queries):
    # Sort rooms by size, then by room id
    rooms.sort(key=lambda x: (x[1], x[0]))
    
    # Add query index for tracking original order
    queries = [(qid, size, i) for i, (qid, size) in enumerate(queries)]
    # Sort queries by size (desc), preferred id, then original index
    queries.sort(key=lambda x: (x[1], x[0], x[2]), reverse=True)
    
    ans = [-1] * len(queries)
    available_rooms = []  # Sorted list of room ids
    
    for qid, size, i in queries:
        # Add all rooms that meet current size requirement
        while rooms and rooms[-1][1] >= size:
            room_id, _ = rooms.pop()
            bisect.insort(available_rooms, room_id)
        
        if available_rooms:
            # Find insertion point for preferred room id
            j = bisect.bisect(available_rooms, qid)
            
            if j == len(available_rooms):
                # Preferred id is larger than all available rooms
                ans[i] = available_rooms[-1]
            elif j == 0:
                # Preferred id is smaller than all available rooms
                ans[i] = available_rooms[0]
            else:
                # Choose closer room (right vs left neighbor)
                if available_rooms[j] - qid < qid - available_rooms[j-1]:
                    ans[i] = available_rooms[j]
                else:
                    ans[i] = available_rooms[j-1]
    
    return ans

# Test the solution
rooms = [[2,2],[1,2],[3,2]]
queries = [[3,1],[3,3],[5,2]]
print(solve(rooms, queries))

[3, -1, 3]

How It Works

The algorithm processes queries in descending order of size requirements. For each query, it adds all eligible rooms (with sufficient size) to a sorted list. Then it uses binary search to find the position where the preferred room id would be inserted, allowing efficient comparison of the closest neighbors.

Time Complexity

The time complexity is O((R + Q) log R) where R is the number of rooms and Q is the number of queries. The space complexity is O(R + Q) for storing the sorted structures.

Conclusion

This solution efficiently finds the closest room for each query by combining sorting with binary search. The key insight is processing queries by decreasing size requirement and maintaining a sorted list of available room ids.