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

1
2
class Solution:
    def closestRoom(self, rooms: List[List[int]], queries: List[List[int]]) -> List[int]:
solution_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
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
hard
参考资料
许可协议

本文除了题目描述（Problem）部分之外，其他内容由 Calf 原创，采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。

分享文章

快来参与讨论吧~