0-1-matrix.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
 https://leetcode.com/problems/01-matrix
"""
__author__ = "prakashsellathurai"
__copyright__ = "Copyright 2021"
__version__ = "1.0.1"
__email__ = "prakashsellathurai@gmail.com"


from collections import deque


class Solution:
    @staticmethod
    def updateMatrix(mat: List[List[int]]) -> List[List[int]]:
        m = len(mat)
        n = len(mat[0])

        dist = [[sys.maxsize for j in range(n)] for i in range(m)]

        for i in range(m):
            for j in range(n):
                if mat[i][j]:
                    if i > 0:
                        dist[i][j] = min(dist[i][j], dist[i - 1][j] + 1)
                    if j > 0:
                        dist[i][j] = min(dist[i][j], dist[i][j - 1] + 1)
                else:
                    dist[i][j] = 0

        for i in range(m - 1, -1, -1):
            for j in range(n - 1, -1, -1):
                if i < m - 1:
                    dist[i][j] = min(dist[i][j], dist[i + 1][j] + 1)
                if j < n - 1:
                    dist[i][j] = min(dist[i][j], dist[i][j + 1] + 1)

        return dist