func updateMatrix(_ mat: [[Int]]) -> [[Int]] {
    // Create an array to hold the result
    var distance = Array(repeating: Array(repeating: Int.max, count: mat[0].count), count: mat.count)
    
    // Iterate through the elements in the matrix
    for row in 0..<mat.count {
        for col in 0..<mat[0].count {
            var queue = [(Int, Int, Int)]() // (row, col, distance)
            // Update the distance if the cell is 0, or add it into the queue if the cell is 1
            if mat[row][col] == 0 {
                distance[row][col] = 0
            } else {
                queue.append((row, col, 1))
            }
            
            // Iterate through the queue to process 1 cells
            while !queue.isEmpty {
                let (currentRow, currentCol, currentDistance) = queue.removeFirst()
                // Move to next 4 neighbour cells to find 1 cell, stop early if the neighbour cell is 0
                for (r, c) in [(-1, 0), (0, 1), (1, 0), (0, -1)] {
                    let (newRow, newCol) = (r + currentRow, c + currentCol)
                    // Ensure newRow and newCol are inside bounds
                    if newRow >= 0 && newRow < mat.count && newCol >= 0 && newCol < mat[0].count {
                        if mat[newRow][newCol] == 0 {
                            distance[row][col] = currentDistance
                            queue.removeAll()
                            break
                        } else {
                            queue.append((newRow, newCol, currentDistance + 1))
                        }
                    }
                }
            }
        }
    }
    
    return distance
}

print(updateMatrix([[0,0,0],[0,1,0],[0,0,0]])) // [[0, 0, 0], [0, 1, 0], [0, 0, 0]]
print(updateMatrix([[0,0,0],[0,1,0],[1,1,1]])) // [[0, 0, 0], [0, 1, 0], [1, 2, 1]]

func updateMatrix(_ mat: [[Int]]) -> [[Int]] {
    var distance = Array(repeating: Array(repeating: Int.max, count: mat[0].count), count: mat.count)
    var queue = [(Int, Int)]() // (row, col)
    for row in 0..<mat.count {
        for col in 0..<mat[0].count {
            if mat[row][col] == 0 {
                queue.append((row, col))
                distance[row][col] = 0
            }
        }
    }

    while !queue.isEmpty {
        let (currentRow, currentCol) = queue.removeFirst()
        for (r, c) in [(-1, 0), (0, 1), (1, 0), (0, -1)] {
            let (newRow, newCol) = (r + currentRow, c + currentCol)
            if newRow >= 0 && newRow < mat.count && newCol >= 0 && newCol < mat[0].count {
                if distance[newRow][newCol] == Int.max {  // this cell has not been visited yet
                    distance[newRow][newCol] = distance[currentRow][currentCol] + 1
                    queue.append((newRow, newCol))
                }
            }
        }
    }

    return distance
}

print(updateMatrix([[0,0,0],[0,1,0],[0,0,0]])) // [[0, 0, 0], [0, 1, 0], [0, 0, 0]]
print(updateMatrix([[0,0,0],[0,1,0],[1,1,1]])) // [[0, 0, 0], [0, 1, 0], [1, 2, 1]]