剑指 Offer II 107. 矩阵中的距离

描述

给定一个由 0 和 1 组成的矩阵 mat ，请输出一个大小相同的矩阵，其中每一个格子是 mat 中对应位置元素到最近的 0 的距离。

两个相邻元素间的距离为 1 。

示例

示例 1:

输入：mat = [[0,0,0],[0,1,0],[0,0,0]]
输出：[[0,0,0],[0,1,0],[0,0,0]]

示例 2:

输入：mat = [[0,0,0],[0,1,0],[1,1,1]]
输出：[[0,0,0],[0,1,0],[1,2,1]]

提示

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 104
• 1 <= m * n <= 104
• mat[i][j] is either 0 or 1.
• mat 中至少有一个 0

解题思路

• 首先找出所有的 0，放入队列
• 然后进行广度优先搜索

代码

class Solution {
    public int[][] updateMatrix(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        boolean[][] seen = new boolean[m][n];
        int[][] ans = new int[m][n];
        Queue<int[]> queue = new LinkedList<int[]>();
        int[] di = {0, 0, -1, 1};
        int[] dj = {1, -1, 0, 0};
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (mat[i][j] == 0) {
                    seen[i][j] = true;
                    queue.offer(new int[]{i ,j});
                }
            }
        }
        while (!queue.isEmpty()) {
            int[] cell = queue.poll();
            int cur_i = cell[0], cur_j = cell[1];
            for (int index = 0; index < 4; index++) {
                int next_i = cur_i + di[index], next_j = cur_j + dj[index];
                if (next_i >= 0 && next_i < m && next_j >= 0 && next_j < n && !seen[next_i][next_j]) {
                    ans[next_i][next_j] = ans[cur_i][cur_j] + 1;
                    seen[next_i][next_j] = true; 
                    queue.offer(new int[]{next_i, next_j});
                }
            }
        }
        return ans;
    }
}