Leetcode 1414 Solution

class Solution:
    def shortestPath(self, grid: List[List[int]], k: int) -> int:
        m = len(grid)
        n = len(grid[0])

        q = collections.deque()
        q.append((0, 0, k))
        visited = {(0, 0, k)}

        ans = 0
        while q:
            s = len(q)

            for _ in range(s):
                i, j, k = q.popleft()

                if i == m - 1 and j == n - 1:
                    return ans

                directions = [
                    (0, 1),
                    (0, -1),
                    (1, 0),
                    (-1, 0),
                ]

                for di, dj in directions:
                    neigh_i = i + di
                    neigh_j = j + dj

                    if not (0 <= neigh_i < m and 0 <= neigh_j < n):
                        continue

                    if grid[neigh_i][neigh_j] == 1:
                        if k > 0 and (neigh_i, neigh_j, k - 1) not in visited:
                            visited.add((neigh_i, neigh_j, k - 1))
                            q.append((neigh_i, neigh_j, k - 1))
                    else:
                        if (neigh_i, neigh_j, k) not in visited:
                            visited.add((neigh_i, neigh_j, k))
                            q.append((neigh_i, neigh_j, k))

            ans += 1

        return -1