542. 01 Matrix

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.

#include<iostream>
#include<vector>
using namespace std;
class Solution {
public:
	vector<vector<int>> updateMatrix(vector<vector<int>>& matrix) {
		vector<vector<int>> dp;
		dp.resize(matrix.size());
		for (int i = 0; i != matrix.size(); ++i) { 
			dp[i].resize(matrix[i].size());
			for (int j = 0; j != dp[i].size(); ++j)
				dp[i][j] = INT_MAX;//初始化, 将dp中每个元素都设为很大的一个数, 这里取INT_MAX, 即距离最近的0的距离为INT_MAX.
		}
		//从dp左上角第一个元素出发, 将dp当前元素与其相邻的4个元素进行比较,判断相邻点与距离1相加后的距离是否比原来的距离要小
		for (int i = 0; i != matrix.size(); ++i){
			for (int j = 0; j != matrix[i].size(); ++j){
				if (matrix[i][j] == 0) dp[i][j] = 0;//matrix中存在元素为0, dp相应位置也为0.
				else{//当前元素与相邻四个元素比较
					if (i > 0 && dp[i-1][j] != INT_MAX && dp[i][j] > dp[i - 1][j] + 1) dp[i][j] = dp[i - 1][j] + 1;
					if (i<matrix.size() - 1 && dp[i + 1][j] != INT_MAX &&  dp[i][j]>dp[i + 1][j] + 1) dp[i][j] = dp[i + 1][j] + 1;
					if (j > 0 && dp[i][j - 1] != INT_MAX && dp[i][j] > dp[i][j - 1] + 1) dp[i][j] = dp[i][j - 1] + 1;
					if (j<matrix[i].size() - 1 && dp[i][j+1] != INT_MAX &&  dp[i][j]>dp[i][j + 1] + 1) dp[i][j] = dp[i][j + 1] + 1;
				}
			}//for
		}//for
		//从dp最后一个元素出发, 反向遍历dp, 找出可能存在的比正向遍历更短的距离
		for (int i = matrix.size()-1; i >=0; --i){
			for (int j = matrix[i].size()-1; j >=0 ; --j){
				if (i > 0 && dp[i - 1][j] != INT_MAX && dp[i][j] > dp[i - 1][j] + 1) dp[i][j] = dp[i - 1][j] + 1;
				if (i<matrix.size() - 1 && dp[i + 1][j] != INT_MAX &&  dp[i][j]>dp[i + 1][j] + 1) dp[i][j] = dp[i + 1][j] + 1;
				if (j > 0 && dp[i][j - 1] != INT_MAX && dp[i][j] > dp[i][j - 1] + 1) dp[i][j] = dp[i][j - 1] + 1;
				if (j<matrix[i].size() - 1 && dp[i][j + 1] != INT_MAX &&  dp[i][j]>dp[i][j + 1] + 1) dp[i][j] = dp[i][j + 1] + 1;
			}//for
		}//for
		return dp;
	}
};
int main(){
	vector<vector<int>> matrix = { { 0, 0, 1, 0, 1, 1, 1, 0, 1, 1 }, { 1, 1, 1, 1, 0, 1, 1, 1, 1, 1 }, { 1, 1, 1, 1, 1, 0, 0, 0, 1, 1}, 
	{ 1, 0, 1, 0, 1, 1, 1, 0, 1, 1}, { 0, 0, 1, 1, 1, 0, 1, 1, 1, 1}, { 1, 0, 1, 1, 1, 1, 1, 1, 1, 1}, { 1, 1, 1, 1, 0, 1, 0, 1, 0, 1}, 
	{ 0, 1, 0, 0, 0, 1, 0, 0, 1, 1}, { 1, 1, 1, 0, 1, 1, 0, 1, 0, 1}, { 1, 0, 1, 1, 1, 0, 1, 1, 1, 0}};
	Solution s;
	vector<vector<int>> disMatrix;
	disMatrix = s.updateMatrix(matrix);
	for (int i = 0; i != disMatrix.size(); ++i){
		for (int j = 0; j != disMatrix[i].size(); ++j)
			cout << disMatrix[i][j] << ' ';
		cout << endl;
	}
	getchar();
}