1847. Closest Room

Problem

There is a hotel with n rooms. The rooms are represented by a 2D integer array rooms where rooms[i] = [roomIdᵢ, sizeᵢ] denotes that there is a room with room number roomIdᵢ and size equal to sizeᵢ. Each roomIdᵢ is guaranteed to be unique.

You are also given k queries in a 2D array queries where queries[j] = [preferredⱼ, minSizeⱼ]. The answer to the jᵗʰ query is the room number id of a room such that:

The room has a size of at least minSizeⱼ, and
abs(id - preferredⱼ) is minimized, where abs(x) is the absolute value of x.

If there is a tie in the absolute difference, then use the room with the smallest such id. If there is no such room, the answer is -1.

Return an array answer of length k where answer[j] contains the answer to the jᵗʰ query.

https://leetcode.cn/problems/closest-room/

Example 1:

Input: rooms = [[2,2],[1,2],[3,2]], queries = [[3,1],[3,3],[5,2]]
Output: [3,-1,3]
Explanation: The answers to the queries are as follows:
Query = [3,1]: Room number 3 is the closest as abs(3 - 3) = 0, and its size of 2 is at least 1. The answer is 3.
Query = [3,3]: There are no rooms with a size of at least 3, so the answer is -1.
Query = [5,2]: Room number 3 is the closest as abs(3 - 5) = 2, and its size of 2 is at least 2. The answer is 3.

Example 2:

Input: rooms = [[1,4],[2,3],[3,5],[4,1],[5,2]], queries = [[2,3],[2,4],[2,5]]
Output: [2,1,3]
Explanation: The answers to the queries are as follows:
Query = [2,3]: Room number 2 is the closest as abs(2 - 2) = 0, and its size of 3 is at least 3. The answer is 2.
Query = [2,4]: Room numbers 1 and 3 both have sizes of at least 4. The answer is 1 since it is smaller.
Query = [2,5]: Room number 3 is the only room with a size of at least 5. The answer is 3.

Constraints:

n == rooms.length
1 <= n <= 10⁵
k == queries.length
1 <= k <= 10⁴
1 <= roomIdᵢ, preferredⱼ <= 10⁷
1 <= sizeᵢ, minSizeⱼ <= 10⁷

Test Cases

python
class Solution:
    def closestRoom(self, rooms: List[List[int]], queries: List[List[int]]) -> List[int]:

solution_test.py下载完整源码
py
import pytest

from solution import Solution
from solution_fast import Solution as SolutionFast
from solution3 import Solution as Solution3


@pytest.mark.parametrize('rooms, queries, expected', [
    ([[2,2],[1,2],[3,2]], [[3,1],[3,3],[5,2]], [3,-1,3]),
    ([[1,4],[2,3],[3,5],[4,1],[5,2]], [[2,3],[2,4],[2,5]], [2,1,3]),

    (
        [[23,22],[6,20],[15,6],[22,19],[2,10],[21,4],[10,18],[16,1],[12,7],[5,22]],
        [[12,5],[15,15],[21,6],[15,1],[23,4],[15,11],[1,24],[3,19],[25,8],[18,6]],
        [12,10,22,15,23,10,-1,5,23,15]
    ),
])
class Test:
    def test_solution(self, rooms, queries, expected):
        sol = Solution()
        assert sol.closestRoom(rooms.copy(), queries.copy()) == expected

    def test_solution_fast(self, rooms, queries, expected):
        sol = SolutionFast()
        assert sol.closestRoom(rooms.copy(), queries.copy()) == expected

    def test_solution3(self, rooms, queries, expected):
        sol = Solution3()
        assert sol.closestRoom(rooms.copy(), queries.copy()) == expected

Thoughts

对于一个 query，可以把所有 size >= minSize 的房间找出来，在这些房间中找距离 preferred 最近的编号。如果这些房间是按照编号排序的，就可以用二分查找。

如果下一个 query 的 minSize 更小一些，那么可以往这个房间集合中再加入那些 size 不小于新 query 的 minSize 的房间，如果还能保持按编号排序，就还可以用二分查找。

所以对 queries 按 minSize 降序排列，这样每下一个 query 的 minSize 都不会比前一个大。对 rooms 也按 size 降序排列，用两个下标分别跟踪 queries 和 rooms，可以方便地对每个新 query 找出增量的房间。

关键在于用一个「能保持有序」的容器保存按 minSize 筛出来的房间，在不断加入新房间的同时保持按房间编号排序。

暂不考虑 grantjenks/python-sortedcontainers: Python Sorted Container Types: Sorted List, Sorted Dict, and Sorted Set 之类的第三方工具，很容易想到的是二叉查找树（Binary Search Tree），或其变体如 AVL 树（Georgy Adelson-Velsky and Evgenii Landis Tree）、红黑树（Red Black Tree）等。

先试了一下 BST，果然很慢（因为非常不平衡）。只好修改插入的方法，使其成为 AVL 树，快了十倍。本题不涉及删除节点，不用红黑树也可以。

AVL 树的详细信息参见 DSA AVL Trees。

稍微修改一下树的查找方法，如果目标数字找不到，就返回小于目标数的最大值、以及大于目标树的最小值，这对于 BST 或其变体都很简单。

时间复杂度是 O(n log n + k log k + k log n)，其中 O(n log n) 是排序 rooms 和不断构建可选房间的有序集合的时间；O(k log k) 是排序 queries 的时间；O(k log n) 是对所有 queries，借助 BST 查找等于或最接近 query 中 preferred 房间编号的时间。

空间复杂度是 O(n + k)，其中 O(n) 是可选房间有序集合的空间（rooms 可以 in-place 排序）；O(k) 是对 queries 排序但需要保留原始下标所需的空间。

提交后的速度不是很快，也就 7+%。可能是 AVL 树的常数系数太大了吧。实际上对于 Python 而言，直接用数组（list）配合标准库里的 bisect — Array bisection algorithm 就可以做到速度非常快（sortedcontainers 底层也是利用 list 和 bisect 库来实现的）。在本题限定的量级内，直接用 list 就已经足够快了（100%）（虽然插入的复杂度是 O(n)）。

Code

AVL Tree

Runtime beats 7+%:

solution.py下载完整源码
py
class TreeNode:
    def __init__(self, val: int, left: 'TreeNode' = None, right: 'TreeNode' = None):
        self.val = val
        self.left = left
        self.right = right
        self.height = 1

    @property
    def balance(self) -> int:
        l = self.left.height if self.left else 0
        r = self.right.height if self.right else 0
        return l - r

    @property
    def left_balance(self) -> int:
        return self.left.balance if self.left else 0

    @property
    def right_balance(self) -> int:
        return self.right.balance if self.right else 0

    def re_height(self) -> None:
        l = self.left.height if self.left else 0
        r = self.right.height if self.right else 0
        self.height = 1 + max(l, r)

    def rotate_left(self) -> 'TreeNode':
        right = self.right
        self.right, right.left = right.left, self
        self.re_height()
        right.re_height()
        return right

    def rotate_right(self) -> 'TreeNode':
        left = self.left
        self.left, left.right = left.right, self
        self.re_height()
        left.re_height()
        return left


class AVLTree:
    def __init__(self) -> None:
        self._root: TreeNode = None

    def insert(self, val: int) -> None:
        if self._root is None:
            self._root = TreeNode(val)
            return

        ancestors = []
        node = self._root
        while True:
            ancestors.append(node)
            if val < node.val:
                if node.left is None:
                    node.left = child = TreeNode(val)
                    break
                else:
                    node = node.left
            else:
                if node.right is None:
                    node.right = child = TreeNode(val)
                    break
                else:
                    node = node.right

        while ancestors:
            node = ancestors.pop()
            if child.val < node.val:
                node.left = child
            else:
                node.right = child
            node.re_height()
            balance = node.balance # left height - right height
            if balance > 1:
                if node.left and node.left.balance < 0: # Left-Right
                    node.left = node.left.rotate_left()
                node = node.rotate_right() # Left-Left & Left-Right
            elif balance < -1:
                if node.right and node.right.balance > 0: # Right-Left
                    node.right = node.right.rotate_right()
                node = node.rotate_left() # # Right-Right & Right-Left
            child = node

        self._root = child

    def find(self, val: int) -> tuple[int, int, int]:
        """Finds `val`. If not found, return the largest value < val, and the smallest value > val."""
        lt = gt = None
        node = self._root
        while node:
            if node.val == val:
                return val, None, None
            elif node.val > val:
                gt = node.val
                node = node.left
            else:
                lt = node.val
                node = node.right

        return None, lt, gt


class Solution:
    def closestRoom(self, rooms: list[list[int]], queries: list[list[int]]) -> list[int]:
        k = len(queries)
        queries = [(q[0], q[1], j) for j, q in enumerate(queries)]
        queries.sort(key=lambda q: q[1], reverse=True)
        answer = [-1] * k

        n = len(rooms)
        rooms.sort(key=lambda r: r[1], reverse=True)
        valid_rooms = AVLTree()

        i = 0
        for preferred, min_size, j in queries:
            while i < n and rooms[i][1] >= min_size:
                valid_rooms.insert(rooms[i][0])
                i += 1

            ans, lt, gt = valid_rooms.find(preferred)
            if ans is None:
                ans = lt
                if gt is not None and (ans is None or gt - preferred < preferred - ans):
                    ans = gt

            if ans is not None: answer[j] = ans

        return answer

Python list - Fast

Runtime beats 100%:

solution_fast.py下载完整源码
py

from bisect import bisect_left, insort


class Solution:
    def closestRoom(self, rooms: list[list[int]], queries: list[list[int]]) -> list[int]:
        k = len(queries)
        queries = [(q[0], q[1], j) for j, q in enumerate(queries)]
        queries.sort(key=lambda q: q[1], reverse=True)
        answer = [-1] * k

        n = len(rooms)
        rooms.sort(key=lambda r: r[1], reverse=True)
        valid_rooms = []

        i = 0
        for preferred, min_size, j in queries:
            while i < n and rooms[i][1] >= min_size:
                insort(valid_rooms, rooms[i][0])
                i += 1

            if valid_rooms:
                idx = bisect_left(valid_rooms, preferred)
                if idx == len(valid_rooms) or idx > 0 and preferred - valid_rooms[idx-1] <= valid_rooms[idx] - preferred:
                    idx -= 1
                answer[j] = valid_rooms[idx]

        return answer

Monotonic Stack

在 2940. Find Building Where Alice and Bob Can Meet 中尝试基于单调栈 + 二分搜索的解法时，觉得这道题也可以用类似的逻辑处理。

核心区别只有两个。一个是本题需要往两个方向找，既要找 preferred 右边第一个，也要找 preferred 左边第一个，可以通过两次循环达成。另一个是 preferred 不一定是存在的房间号，不能用跟 problem 2940 一样的方式（即用当前扫描到的房间号查出 prefer 此房间号的所有查询），需要对 preferred 和 roomId 做分段匹配。

各种边界条件细节需要小心处理。

时间复杂度是 O(n log n + k log k + k log n)，空间复杂度 O(n + k)，跟上边一样。实际运行时间跟 problem 2940 也很像，用单调栈比用有序集合慢一倍。

solution3.py下载完整源码
py

from bisect import bisect_right

gt = lambda a, b: a > b
lt = lambda a, b: a < b


class Solution:
    def closestRoom(self, rooms: list[list[int]], queries: list[list[int]]) -> list[int]:
        rooms.sort() # Sort rooms by roomId.
        queries = [(q[0], q[1], j) for j, q in enumerate(queries)] # Remember origin query index before sort.
        queries.sort() # Sort queries by preferred.
        answer = [-1] * len(queries)

        for done, valid in enumerate((gt, lt)):
            j = len(queries) - 1
            while j >= 0 and valid(queries[j][0], rooms[-1][0]): # Skip no-result queries for current direction.
                j -= 1

            stack = []
            for i in range(len(rooms) - 1, -1, -1):
                size = rooms[i][1]
                while stack and rooms[stack[-1]][1] <= size:
                    stack.pop()
                stack.append(i)

                while j >= 0 and (i == 0 or valid(queries[j][0], rooms[i-1][0])):
                    preferred, min_size, qid = queries[j]
                    j -= 1
                    if answer[qid] == preferred: continue

                    idx = bisect_right(stack, -min_size, key=lambda x: -rooms[x][1]) - 1
                    if idx < 0: continue
                    if answer[qid] == -1 or preferred - rooms[stack[idx]][0] <= answer[qid] - preferred:
                        answer[qid] = rooms[stack[idx]][0]

            if done: break
            rooms.reverse()
            queries.reverse()

        return answer

#Problem

#Test Cases

#Thoughts

#Code

=AVL Tree

=Python list - Fast

#Monotonic Stack