BFS and solve the problem reversely

542. 01 Matrix

If we solve it naively based on the problem description, we will get TLE.

const int INF = 0x3f3f3f3f;
vector<pair<int, int>> dirs = {{0, 1},{0, -1},{1, 0},{-1, 0}};
class Solution {
public:
    void bfs(int i, int j, vector<vector<int>>& mat,  vector<vector<int>>& dis){
        vector<pair<int, int>> Q;
        Q.emplace_back(i, j);
        const int n = mat.size(), m = mat[0].size();
        vector<vector<int>> seen(n, vector<int>(m, 0));
        int level = 1;
        seen[i][j] = 1;
        int currDis = INF;
        while (!Q.empty()){
            vector<pair<int, int>> newQ;
            for (auto [ii, jj]: Q){
                for (auto [di, dj]: dirs){
                    int r = ii + di, c = jj + dj;
                    if (r < 0 || r >= n || c < 0 || c >= m) continue;
                    if (mat[r][c] == 0) {
                        dis[i][j] = level;
                        return;
                    } else if (!seen[r][c]){
                        if (dis[r][c] != INF && dis[r][c] + level > currDis)continue;
                        if (dis[r][c] != INF) currDis = min(currDis, level + dis[r][c]);
                        seen[r][c] = 1;
                        newQ.emplace_back(r, c);
                    }
                }
            }
            Q = newQ;
            level++;
        }
        dis[i][j] = currDis;
        return;
    }
    vector<vector<int>> updateMatrix(vector<vector<int>>& mat) {
        const int n = mat.size(), m = mat[0].size();
        vector<vector<int>> dis(n, vector<int>(m, INF));
        for (int i = 0; i < n; ++i){
            for (int j = 0; j < m; ++j){
                if (!mat[i][j]) dis[i][j] = 0;
            }
        }
        for (int i = 0; i < n; ++i){
            for (int j = 0; j < m; ++j){
                if (mat[i][j]) bfs(i, j, mat, dis);
            }
        }
        return dis;
        
    }
};

Got TLE

37 / 49 test cases passed.

If we think it reversely by exploring from 0 to non 0 cells, then we can solve this problem efficiently by using BFS.

vector<pair<int, int>> dirs = {{0, 1},{0, -1},{1, 0},{-1, 0}};
class Solution {
public:
    vector<vector<int>> updateMatrix(vector<vector<int>>& mat) {
        vector<pair<int, int>> Q;
        const int n = mat.size(), m = mat[0].size();
        vector<vector<int>> dis(n, vector<int>(m));
        vector<vector<int>> seen(n, vector<int>(m, 0));
        for (int i = 0; i < n; ++i){
            for (int j = 0; j < m; ++j){
                if (!mat[i][j]) {
                    Q.emplace_back(i, j);
                    seen[i][j] = 1;dis[i][j] = 0;
                }
            }
        }
        for (auto [i, j]: Q)seen[i][j] = 1;
        
        // BFS
        int level = 1;
        while (!Q.empty()){
            vector<pair<int, int>> newQ;
            for (auto [ii, jj]: Q){
                for (auto [di, dj]: dirs){
                    int r = ii + di, c = jj + dj;
                    if (r < 0 || r >= n || c < 0 || c >= m || seen[r][c]) continue;
                    seen[r][c] = 1; dis[r][c] = level;
                    newQ.emplace_back(r, c);
                }
            }
            Q = newQ;level++;
        }    
        return dis;
        
    }
};