LeetCode in Kotlin

1847. Closest Room

Hard

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

You are also given k queries in a 2D array queries where queries[j] = [preferredj, minSizej]. The answer to the jth query is the room number id of a room such that:

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 jth query.

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:

Solution

import java.util.TreeSet

class Solution {
    fun closestRoom(rooms: Array<IntArray>, queries: Array<IntArray>): IntArray {
        val numRoom = rooms.size
        val numQuery = queries.size
        for (i in 0 until numQuery) {
            queries[i] = intArrayOf(queries[i][0], queries[i][1], i)
        }
        rooms.sortWith { a: IntArray, b: IntArray -> if (a[1] != b[1]) a[1] - b[1] else a[0] - b[0] }
        queries.sortWith { a: IntArray, b: IntArray -> if (a[1] != b[1]) a[1] - b[1] else a[0] - b[0] }
        val roomIds = TreeSet<Int>()
        val result = IntArray(numQuery)
        var j = numRoom - 1
        for (i in numQuery - 1 downTo 0) {
            val currRoomId = queries[i][0]
            val currRoomSize = queries[i][1]
            val currQueryIndex = queries[i][2]
            while (j >= 0 && rooms[j][1] >= currRoomSize) {
                roomIds.add(rooms[j--][0])
            }
            if (roomIds.contains(currRoomId)) {
                result[currQueryIndex] = currRoomId
                continue
            }
            val nextRoomId = roomIds.higher(currRoomId)
            val prevRoomId = roomIds.lower(currRoomId)
            if (nextRoomId == null && prevRoomId == null) {
                result[currQueryIndex] = -1
            } else if (nextRoomId == null) {
                result[currQueryIndex] = prevRoomId!!
            } else if (prevRoomId == null) {
                result[currQueryIndex] = nextRoomId
            } else {
                if (currRoomId - prevRoomId <= nextRoomId - currRoomId) {
                    result[currQueryIndex] = prevRoomId
                } else {
                    result[currQueryIndex] = nextRoomId
                }
            }
        }
        return result
    }
}