/src/main/java/org/rajivprab/sandbox/submatrix_target/SubmatrixSumTargetOptimized.java

package org.rajivprab.sandbox.submatrix_target;

import com.google.common.collect.HashBasedTable;
import com.google.common.collect.ImmutableTable;
import com.google.common.collect.Maps;
import com.google.common.collect.Table;

import java.util.Map;
import java.util.Optional;

/**
 * See: https://leetcode.com/problems/number-of-submatrices-that-sum-to-target/discuss/303750/JavaC%2B%2BPython-Find-the-Subarray-with-Target-Sum
 *
 * Complexity: O(Row*Col*Col)
 */
public class SubmatrixSumTargetOptimized {

    // Matrix.get(row, column) = value
    public static int getNumSubmatrices(Table<Integer, Integer, Integer> matrix, int target, int numRows, int numCol) {
        return new SubmatrixSumTargetOptimized(matrix, numRows, numCol).getNumSubmatricesForTarget(target);
    }

    private final Table<Integer, Integer, Integer> matrix;
    private final Table<Integer, Integer, Integer> rowSumMatrix;
    private final int numRows;
    private final int numColumns;

    private SubmatrixSumTargetOptimized(Table<Integer, Integer, Integer> matrix, int numRows, int numCol) {
        this.matrix = ImmutableTable.copyOf(matrix);
        this.numRows = numRows;
        this.numColumns = numCol;
        this.rowSumMatrix = ImmutableTable.copyOf(buildRowSumMatrix());
    }

    private Table<Integer, Integer, Integer> buildRowSumMatrix() {
        Table<Integer, Integer, Integer> rowSumMatrix = HashBasedTable.create();
        for (int row=0; row<numRows; row++) {
            for (int col=0; col<numColumns; col++) {
                int left = Optional.ofNullable(rowSumMatrix.get(row, col-1)).orElse(0);
                int cur = left + matrix.get(row, col);
                rowSumMatrix.put(row, col, cur);
            }
        }
        return rowSumMatrix;
    }

    private int getNumSubmatricesForTarget(int target) {
        int numResults = 0;
        for (int startColumn = 0; startColumn < numColumns; startColumn++) {
            for (int endColumn = startColumn; endColumn < numColumns; endColumn++) {
                numResults += getNumSubmatricesForTarget(startColumn, endColumn, target);
            }
        }
        return numResults;
    }

    private int getNumSubmatricesForTarget(int startColumn, int endColumn, int target) {
        int numResults = 0;

        Map<Integer, Integer> sumCounts = Maps.newHashMap();
        int curSum = 0;
        sumCounts.put(curSum, 1);
        for (int row=0; row<numRows; row++) {
            curSum += getRowSum(row, startColumn, endColumn);

            int startSum = curSum - target;
            numResults += sumCounts.getOrDefault(startSum, 0);

            int updatedCount = sumCounts.getOrDefault(curSum, 0) + 1;
            sumCounts.put(curSum, updatedCount);
        }
        return numResults;
    }

    private int getRowSum(int row, int startColumn, int endColumn) {
        return rowSumMatrix.get(row, endColumn) -
                Optional.ofNullable(rowSumMatrix.get(row, startColumn - 1)).orElse(0);
    }
}