LeetCode in Kotlin

2503. Maximum Number of Points From Grid Queries

Hard

You are given an m x n integer matrix grid and an array queries of size k.

Find an array answer of size k such that for each integer queries[i] you start in the top left cell of the matrix and repeat the following process:

After the process, answer[i] is the maximum number of points you can get. Note that for each query you are allowed to visit the same cell multiple times.

Return the resulting array answer.

Example 1:

Input: grid = [[1,2,3],[2,5,7],[3,5,1]], queries = [5,6,2]

Output: [5,8,1]

Explanation: The diagrams above show which cells we visit to get points for each query.

Example 2:

Input: grid = [[5,2,1],[1,1,2]], queries = [3]

Output: [0]

Explanation: We can not get any points because the value of the top left cell is already greater than or equal to 3.

Constraints:

Solution

import java.util.ArrayDeque
import java.util.Arrays
import java.util.PriorityQueue
import java.util.Queue

class Solution {
    private val dirs = arrayOf(intArrayOf(-1, 0), intArrayOf(1, 0), intArrayOf(0, -1), intArrayOf(0, 1))

    fun maxPoints(grid: Array<IntArray>, queries: IntArray): IntArray {
        val r = grid.size
        val c = grid[0].size
        val res = IntArray(queries.size)
        val index = arrayOfNulls<Int>(queries.size)
        for (i in queries.indices) {
            index[i] = i
        }
        Arrays.sort(index, { o: Int?, m: Int? -> queries[o!!].compareTo(queries[m!!]) })
        val q1: Queue<IntArray> = ArrayDeque()
        val q2 = PriorityQueue({ a: IntArray, b: IntArray -> a[2].compareTo(b[2]) })
        q2.offer(intArrayOf(0, 0, grid[0][0]))
        val visited = Array(r) { BooleanArray(c) }
        var count = 0
        visited[0][0] = true
        for (i in queries.indices) {
            val currLimit = queries[index[i]!!]
            while (q2.isNotEmpty() && q2.peek()[2] < currLimit) {
                q1.offer(q2.poll())
            }
            while (q1.isNotEmpty()) {
                val curr = q1.poll()
                count++
                for (dir in dirs) {
                    val x = dir[0] + curr[0]
                    val y = dir[1] + curr[1]
                    if (x < 0 || y < 0 || x >= r || y >= c || visited[x][y]) {
                        continue
                    }
                    if (!visited[x][y] && grid[x][y] < currLimit) {
                        q1.offer(intArrayOf(x, y, grid[x][y]))
                    } else {
                        q2.offer(intArrayOf(x, y, grid[x][y]))
                    }
                    visited[x][y] = true
                }
            }
            res[index[i]!!] = count
        }
        return res
    }
}