def print_matrix(matrix):
    row_num = 0
    print("\t\t      [", end = "")

    for i in range(len(matrix)):

        if row_num == 0:
            print("[", end="")
        else:
            print("\t\t       [", end="")


        for j in range(len(matrix[i])):
            print(matrix[i][j], end="")
            if j != len(matrix[i]) - 1:
                print(", ", end="")
        
        
        if row_num != len(matrix) - 1:
            print("]", end="")
        else:
            print("]]", end="")
        row_num += 1
        if i != len(matrix) - 1:
            print(",", end=" ")
        print()

import math

def update_matrix(mat):
    # calculate the number of rows and columns in matrix
    m, n = len(mat), len(mat[0])

    for r in range(m):
        for c in range(n):

            # check if the element is greater than zero
            if mat[r][c] > 0:

                # check the element above, if there is no element above, set to infinity
                above = mat[r - 1][c] if r > 0 else math.inf

                # check the left element, if there is no left element, set to infinity
                left = mat[r][c - 1] if c > 0 else math.inf

                # update the current element with the minimum of element above and to its left, + 1
                mat[r][c] = min(above, left) + 1

    for r in range(m - 1, -1, -1):
        for c in range(n - 1, -1, -1):

            # check if the element is greater than zero
            if mat[r][c] > 0:

                # check the element below, if there is no element below, set to infinity
                below = mat[r + 1][c] if r < m - 1 else math.inf

                # check the right element, if there is no right element, set to infinity
                right = mat[r][c + 1] if c < n - 1 else math.inf

                # update the current element with the minimum of its value, element below and to its right, + 1
                mat[r][c] = min(mat[r][c], below + 1, right + 1)
    return mat

# Driver code
def main():
    input_bits = [
        [[0, 1], [1, 1]],
        [[0, 0, 1], [0, 1, 1], [1, 0, 1]],
        [[0, 0, 0], [0, 1, 0], [1, 0, 1]],
        [[0, 0, 0], [0, 1, 0], [1, 1, 1]],
        [[0, 1, 0, 1], [1, 1, 1, 0], [0, 1, 1, 1], [1, 0, 1, 1]],
    ]

    for i in range(len(input_bits)):
        print(i + 1, ".\t Input matrix: ", sep="")
        print_matrix(input_bits[i])
        print("\n\t Distance matrix: ", sep="")
        print_matrix(update_matrix(input_bits[i]))
        print("-" * 100)

if __name__ == "__main__":
    main()