Lintcode 405. Submatrix Sum

Given an integer matrix, find a submatrix where the sum of numbers is zero. Your code should return the coordinate of the left-up and right-down number.

If there are multiple answers, you can return any of them.

Example

Example 1:

Input:
[
  [1, 5, 7],
  [3, 7, -8],
  [4, -8 ,9]
]
Output: [[1, 1], [2, 2]]

Example 2:

Input:
[
  [0, 1],
  [1, 0]
]
Output: [[0, 0], [0, 0]]

Challenge

O(n3) time.

Naive Solution:
A naive solution is first calculate all the presum[x][y], which means all presums from [0][0] to [x][y]. Then for each two pairs (x1, y1) and (x2, y2) we can calculate the sum of the rectangle in O(1) time. So the overall time complexity is O(n^4)

A better Solution:
A better solution is we first calculate the presum for each column, and then we can enumerate each two rows, say row1 and row2, we can easily calculate the presum of each column between [row1, row2]. Then the problem can be downgraded to a 1D problem. 

Code (Java):

public class Solution {
    /*
     * @param matrix: an integer matrix
     * @return: the coordinate of the left-up and right-down number
     */
    public int[][] submatrixSum(int[][] matrix) {
        int[][] ans = new int[2][];
        for (int i = 0; i < 2; i++) {
            ans[i] = new int[2];
        }
        
        if (matrix == null || matrix.length == 0) {
            return ans;
        }
        
        int m = matrix.length;
        int n = matrix[0].length;
        int[][] preSum = new int[m + 1][n + 1];
        
        for (int i = 1; i <= m; i++) {
            for (int j = 1; j <= n; j++) {
                preSum[i][j] = preSum[i - 1][j] + matrix[i - 1][j - 1];
            }
        }
        
        for (int x1 = 1; x1 <= m; x1++) {
            for (int x2 = x1; x2 <= m; x2++) {
                Map<Integer, Integer> map = new HashMap<>();
                map.put(0, 0);
                int sum = 0;
                for (int y2 = 1; y2 <= n; y2++) {
                    sum += preSum[x2][y2] - preSum[x1 - 1][y2];
                    if (map.containsKey(sum)) {
                        ans[0][0] = x1 - 1;
                        ans[0][1] = map.get(sum);
                        ans[1][0] = x2 - 1;
                        ans[1][1] = y2 - 1;
                        
                        return ans;
                    }
                    map.put(sum, y2);
                }
            }
        }
        
        return ans;
    }
}