Given a non-empty 2D matrix matrix and an integer k, find the max sum of a rectangle in the matrix such that its sum is no larger than k.

Example:
Given matrix = [
  [1,  0, 1],
  [0, -2, 3]
]
k = 2
The answer is 2. Because the sum of rectangle [[0, 1], [-2, 3]] is 2 and 2 is the max number no larger than k (k = 2).

Note: The rectangle inside the matrix must have an area > 0. What if the number of rows is much larger than the number of columns?

Solution

public class Solution {
    public int maxSumSubmatrix(int[][] matrix, int K) {
        if(matrix.length == 0 || matrix[0].length == 0) return 0;

        // the key is how to get the array no larger than k
        Integer ans = null;
        int m = matrix.length;
        int n = matrix[0].length;
        int[] buffer = new int[m];
        for(int i = 0; i < n; i++) {
            Arrays.fill(buffer, 0);
            for(int j = i; j < n; j++) {
                for(int k = 0; k < m; k++) {
                    buffer[k] += matrix[k][j];
                }
                Integer max = getMaxSumK(buffer, K);
                if(max != null && (ans == null || max > ans)) ans = max;
            }
        }
        return ans == null ? 0 : ans;
    }

    public Integer getMaxSumK(int[] buffer, int k) {
        Integer ans = null;

        long sum[] = new long[buffer.length+1];
        for(int i = 0; i < buffer.length; i++) sum[i+1] = sum[i] + buffer[i];

        TreeSet<Long> tree = new TreeSet<Long>();
        tree.add(0L);
        for(int i = 1; i < sum.length; i++) {
            Long si = tree.ceiling(sum[i] - k);
            if(si != null && (ans == null || (sum[i]-si) > ans)) {
                ans = (int)(sum[i] - si);
            }
            tree.add(sum[i]);
        }


        /*
        for(int i = 1; i < sum.length; i++) {
            for(int j = 0; j < i; j++) {
                long tmp = sum[i] - sum[j];
                if(tmp <= k && (ans == null || tmp > ans)) {
                    ans = (int)tmp;
                }
            }
        }*/

        return ans; 
    }
}