542. 01 Matrix (Medium)

Given a matrix consists of 0 and 1, find the distance of the nearest 0 for each cell.

The distance between two adjacent cells is 1.

Example 1:

Input:

0 0 0
0 1 0
0 0 0

Output:

0 0 0
0 1 0
0 0 0

Example 2: Input:

0 0 0
0 1 0
1 1 1

Output:

0 0 0
0 1 0
1 2 1

Note:

1. The number of elements of the given matrix will not exceed 10,000.
2. There are at least one 0 in the given matrix.
3. The cells are adjacent in only four directions: up, down, left and right.

Solution: 218ms

class Solution {
public:
    vector<vector<int>> updateMatrix(vector<vector<int>>& matrix) {
        int m = matrix.size(), n = matrix[0].size();
        queue<int> q;

        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                if (matrix[i][j] == 0) q.push(i*n+j);
                else matrix[i][j] = INT_MAX;
            }
        }

        vector<int> dirs = {0,-1,0,1,0};
        while (!q.empty()) {
            int r = q.front()/n, c = q.front()%n; q.pop();
            int val = matrix[r][c];
            for (int j = 0; j < 4; ++j) {
                int x = r+dirs[j], y = c+dirs[j+1], idx = x*n+y;
                if (x < 0 || x >= m || y < 0 || y >= n || 
                    matrix[x][y] <= val+1) continue;
                matrix[x][y] = val+1;
                q.push(idx);
            }
        }
        return matrix;
    }
};