数据结构-全部由1组成的子矩形数量

给定一个二维数组matrix, 其中的值不是0就是1， 返回全部由1组成的子矩形数量。

import java.util.Stack;

public class CountSubmatricesWithAllOnes {
    public static void main(String[] args) {
        int[][] mat = {
                {1,1,1,1,1,1},
                {1,1,1,1,1,1},
                {1,1,1,1,1,1}
        };

        System.out.println(numSubmat(mat));
    }

    public static int numSubmat(int[][] mat){
        if(mat == null || mat.length == 0 || mat[0].length == 0){
            return 0;
        }

        int nums = 0;
        int[] height = new int[mat[0].length];
        for (int i = 0; i < mat.length; i++) {
            for (int j = 0; j < mat[0].length; j++) {
                height[j] = mat[i][j] == 0 ? 0 : (height[j] + 1);
            }
            nums += countFromBottom(height);
        }

        return nums;
    }

    public static int countFromBottom(int[] height){
        if(height == null || height.length == 0){
            return 0;
        }

        int nums = 0;
        Stack<Integer> stack = new Stack<Integer>();
        for (int i = 0; i < height.length; i++) {
            if(!stack.isEmpty() && height[stack.peek()] >= height[i]){
                if(height[stack.peek()] == height[i]){
                    stack.pop(); // 如果相等，就弹出栈，不计算当天弹出的索引
                    stack.push(i);
                    continue;
                }

                int popIndex = stack.pop();
                int h = height[popIndex];
                int leftIndex = stack.isEmpty() ? -1 : stack.peek();
                int lenght = i - leftIndex - 1; // 长度
                int leftHight = leftIndex == - 1 ? 0 : height[leftIndex];
                int rightHight = height[i];
                nums = nums + ( lenght * (lenght + 1) / 2 *  (h - Math.max(leftHight,rightHight)) );
            }

            stack.push(i);
        }

        while(!stack.isEmpty()){
            int popIndex = stack.pop();
            int h = height[popIndex];
            int leftIndex = stack.isEmpty() ? -1 : stack.peek();
            int lenght = height.length - leftIndex - 1; // 长度
            int leftHight = leftIndex == - 1 ? 0 : height[leftIndex];
            int rightHight = 0;
            nums = nums + ( lenght * (lenght + 1) / 2 *  (h - Math.max(leftHight,rightHight)) );
        }

        return nums;
    }
}