Lintcode 631. Maximal Square II

Given a 2D binary matrix filled with 0's and 1's, find the largest square which diagonal is all 1 and others is 0.

Example

Example 1:

Input:
[[1,0,1,0,0],[1,0,0,1,0],[1,1,0,0,1],[1,0,0,1,0]]
Output:
9
Explanation:
[0,2]->[2,4]

Example 2:

Input:
[[1,0,1,0,1],[1,0,0,1,1],[1,1,1,1,1],[1,0,0,1,0]]
Output:
4
Explanation:
[0,2]->[1,3]

Notice

Only consider the main diagonal situation.

Analysis:
A DP problem. We define dp[i][j] the max length of square which diagonal is all 1 and others is 0. 

So dp[i][j] depends on three things:
a. dp[i - 1][j - 1], 
b. the number of contiguous 0s starting from matrix[i][j - 1] to left 
c. the number of contiguous 0s starting from matrix[i - 1][j] to top

the dp[i][j] = min(dp[i - 1][j - 1], a, b) + 1;

Code (Java):

public class Solution {
    /**
     * @param matrix: a matrix of 0 an 1
     * @return: an integer
     */
    public int maxSquare2(int[][] matrix) {
        if (matrix == null || matrix.length == 0) {
            return 0;
        }

        int maxLen = 0;

        int m = matrix.length;
        int n = matrix[0].length;

        int[][] rowZeros = new int[m][n]; // most contignous zeros 
        int[][] colZeros = new int[m][n];
        int[][] dp = new int[2][n];

        for (int i = 0 ; i< m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1) {
                    rowZeros[i][j] = 0;
                    colZeros[i][j] = 0;
                } else {
                    rowZeros[i][j] = j == 0 ? 1 : rowZeros[i][j - 1] + 1;
                    colZeros[i][j] = i == 0 ? 1 : colZeros[i - 1][j] + 1;
                }
            }
        }

        for (int i = 0; i < n; i++) {
            dp[0][i] = matrix[0][i];
            maxLen = Math.max(maxLen, dp[0][i]);
        }

        int cur = 0;
        int old = 0;
        for (int i = 1; i < m; i++) {
            old = cur;
            cur = 1 - cur;
            dp[cur][0] = matrix[i][0];
            maxLen = Math.max(maxLen, dp[cur][0]);
            for (int j = 1; j < n; j++) {
                if (matrix[i][j] == 0) {
                    dp[cur][j] = 0;
                } else {
                    dp[cur][j] = Math.min(dp[old][j - 1], Math.min(rowZeros[i][j - 1], colZeros[i - 1][j])) + 1;
                }

                maxLen = Math.max(maxLen, dp[cur][j]);
            }
        }

        return maxLen * maxLen;
    }
}