Skip to content

1847. Closest Room 👍

  • Time: $O(\texttt{sort}(n) + \texttt{sort}(q) + q\log n)$
  • Space: $O(n + q)$
 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
class Solution {
 public:
  vector<int> closestRoom(vector<vector<int>>& rooms,
                          vector<vector<int>>& queries) {
    vector<int> ans(queries.size());
    set<int> roomIds;

    for (int i = 0; i < queries.size(); ++i)
      queries[i].push_back(i);

    ranges::sort(rooms, ranges::greater{},
                 [](const vector<int>& room) { return room[1]; });
    ranges::sort(queries, ranges::greater{},
                 [](const vector<int>& query) { return query[1]; });

    int i = 0;  // rooms' index
    for (const vector<int>& query : queries) {
      while (i < rooms.size() && rooms[i][1] >= query[1])
        roomIds.insert(rooms[i++][0]);
      ans[query[2]] = searchClosestRoomId(roomIds, query[0]);
    }

    return ans;
  }

 private:
  int searchClosestRoomId(set<int>& roomIds, int preferred) {
    const auto it = roomIds.lower_bound(preferred);
    const int id1 = it == roomIds.cbegin() ? -1 : *(prev(it));
    const int id2 = it == roomIds.cend() ? -1 : *it;
    if (id1 == -1)
      return id2;
    if (id2 == -1)
      return id1;
    if (abs(preferred - id1) <= abs(preferred - id2))
      return id1;
    return id2;
  }
};

 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
class Solution {
  public int[] closestRoom(int[][] rooms, int[][] queries) {
    int[] ans = new int[queries.length];
    Integer[] indices = new Integer[queries.length];
    TreeSet<Integer> roomIds = new TreeSet<>();

    for (int i = 0; i < queries.length; ++i)
      indices[i] = i;

    Arrays.sort(rooms, Comparator.comparingInt(room -> - room[1]));
    Arrays.sort(indices, Comparator.comparingInt(index -> - queries[index][1]));

    int i = 0; // rooms' index
    for (final int index : indices) {
      while (i < rooms.length && rooms[i][1] >= queries[index][1])
        roomIds.add(rooms[i++][0]);
      ans[index] = searchClosestRoomId(roomIds, queries[index][0]);
    }

    return ans;
  }

  private int searchClosestRoomId(TreeSet<Integer> roomIds, int preferred) {
    Integer floor = roomIds.floor(preferred);
    Integer ceiling = roomIds.ceiling(preferred);
    final int id1 = floor == null ? -1 : floor;
    final int id2 = ceiling == null ? -1 : ceiling;
    if (id1 == -1)
      return id2;
    if (id2 == -1)
      return id1;
    if (Math.abs(preferred - id1) <= Math.abs(preferred - id2))
      return id1;
    return id2;
  }
}

 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
from sortedcontainers import SortedList


class Solution:
  def closestRoom(
      self,
      rooms: list[list[int]],
      queries: list[list[int]],
  ) -> list[int]:
    ans = [0] * len(queries)
    qs = [[*q, i] for i, q in enumerate(queries)]
    roomIds = SortedList()

    rooms.sort(key=lambda x: -x[1])
    qs.sort(key=lambda x: -x[1])

    def searchClosestRoomId(roomIds: SortedList, preferred: int):
      if not roomIds:
        return -1

      candIds = []
      i = roomIds.bisect_right(preferred)
      if i > 0:
        candIds.append(roomIds[i - 1])
      if i < len(roomIds):
        candIds.append(roomIds[i])
      return min(candIds, key=lambda x: abs(x - preferred))

    i = 0  # rooms' index
    for preferred, minSize, index in qs:
      while i < len(rooms) and rooms[i][1] >= minSize:
        roomIds.add(rooms[i][0])
        i += 1
      ans[index] = searchClosestRoomId(roomIds, preferred)

    return ans