LeetCode - Count Square Submatrices with All Ones

Given a m * n matrix of ones and zeros, return how many square submatrices have all ones.

Example 1:

Input: matrix = [ [0,1,1,1], [1,1,1,1], [0,1,1,1] ] Output: 15 Explanation: There are 10 squares of side 1. There are 4 squares of side 2. There is 1 square of side 3. Total number of squares = 10 + 4 + 1 = 15.

Example 2:

Input: matrix = [ [1,0,1], [1,1,0], [1,1,0] ] Output: 7 Explanation: There are 6 squares of side 1. There is 1 square of side 2. Total number of squares = 6 + 1 = 7.

Constraints:

1 <= arr.length <= 300
1 <= arr[0].length <= 300
0 <= arr[i][j] <= 1

Solution

Its a matrix DP problem. Every location in the matrix can have two values 0/1, If its a 0 then it can not contribute to the total squares count, if its 1 then it can.

If location having value as 1, then element itself can be consider as square, Also it can be part of other square ending at the current location, that can be calculated based on previous row and column having 1s or not.

For example

Input: matrix = [ [1,0], [1,1] ] Explanation : At location (2,2) value is 1, But it can not contribute to the other square because at location (0,1) having value as 0.

DP equation for that can be formed as Min(Matrix[i-1][j-1], Matrix[i-1][j], Matrix[i][j-1])+1

Code


public class CountSquareMatrix {


	public int countSquares(int[][] matrix) {
		
		int cnt = 0;

		if(matrix.length==0) return cnt;
		
		/**
		 * Checking first column for all ones. 
		 */
		for (int i = 0; i < matrix.length; i++) {
			if (matrix[i][0] == 1) {
				cnt++;
			}
		}
		
		/**
		 * Checking first row for all ones. 
		 */
		for (int i = 1; i < matrix[0].length; i++) {
			if (matrix[0][i] == 1) {
				cnt++;
			}
		}
		
		for (int i = 1; i < matrix.length; i++) {

			for (int j = 1; j < matrix[i].length; j++) {

				if (matrix[i][j] == 1) {
					matrix[i][j] = Math.min(matrix[i-1][j-1],Math.min(matrix[i - 1][j],matrix[i][j - 1])) + 1;
					cnt += matrix[i][j];
				}
			}
		}
		return cnt;

	}
	public static void main(String[] args) {

		int[][] matrix =  {
				{1,0,1},
				{1,1,0},
				{1,1,0}};
		
		CountSquareMatrix csMatrix=new CountSquareMatrix();
		System.out.println(csMatrix.countSquares(matrix));

	}
}