Range Sum Query 2D

Author: nlohia-godaddy

src/main/java/leetcode/medium/RangeSumQuery2D.java

@@ -0,0 +1,37 @@
+package leetcode.medium;
+
+public class RangeSumQuery2D {
+
+  private int[][] prefix;
+
+  public RangeSumQuery2D(int[][] matrix) {
+    if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
+      return;
+    }
+
+    int m = matrix.length;
+    int n = matrix[0].length;
+
+    prefix = new int[m][n];
+
+    for (int i = 0; i < m; i++) {
+      for (int j = 0; j < n; j++) {
+        int top = (i > 0) ? prefix[i - 1][j] : 0;
+        int left = (j > 0) ? prefix[i][j - 1] : 0;
+        int topLeft = (i > 0 && j > 0) ? prefix[i - 1][j - 1] : 0;
+
+        prefix[i][j] = matrix[i][j] + top + left - topLeft;
+      }
+    }
+  }
+
+  public int sumRegion(int row1, int col1, int row2, int col2) {
+    int total = prefix[row2][col2];
+    int top = (row1 > 0) ? prefix[row1 - 1][col2] : 0;
+    int left = (col1 > 0) ? prefix[row2][col1 - 1] : 0;
+    int topLeft = (row1 > 0 && col1 > 0) ? prefix[row1 - 1][col1 - 1] : 0;
+
+    return total - top - left + topLeft;
+  }
+
+}

src/test/java/leetcode/medium/RangeSumQuery2DTest.java

@@ -0,0 +1,60 @@
+package leetcode.medium;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class RangeSumQuery2DTest {
+
+  @Test
+  void testBasicSum() {
+    int[][] matrix = {
+        {1, 2, 3},
+        {4, 5, 6},
+        {7, 8, 9}
+    };
+    RangeSumQuery2D rsq = new RangeSumQuery2D(matrix);
+    assertEquals(12, rsq.sumRegion(0, 0, 1, 1)); // 1+2+4+5
+  }
+
+  @org.junit.jupiter.api.Test
+  void testEntireMatrix() {
+    int[][] matrix = {
+        {1, 2},
+        {3, 4}
+    };
+    RangeSumQuery2D rsq = new RangeSumQuery2D(matrix);
+    assertEquals(10, rsq.sumRegion(0, 0, 1, 1)); // 1+2+3+4
+  }
+
+  @org.junit.jupiter.api.Test
+  void testSingleElement() {
+    int[][] matrix = {
+        {5, 6},
+        {7, 8}
+    };
+    RangeSumQuery2D rsq = new RangeSumQuery2D(matrix);
+    assertEquals(6, rsq.sumRegion(0, 1, 0, 1)); // 6
+  }
+
+  @org.junit.jupiter.api.Test
+  void testRowQuery() {
+    int[][] matrix = {
+        {1, 2, 3},
+        {4, 5, 6}
+    };
+    RangeSumQuery2D rsq = new RangeSumQuery2D(matrix);
+    assertEquals(6, rsq.sumRegion(0, 0, 0, 2)); // 1+2+3
+  }
+
+  @org.junit.jupiter.api.Test
+  void testColumnQuery() {
+    int[][] matrix = {
+        {1, 2},
+        {3, 4},
+        {5, 6}
+    };
+    RangeSumQuery2D rsq = new RangeSumQuery2D(matrix);
+    assertEquals(9, rsq.sumRegion(0, 0, 2, 0)); // 1+3+5
+  }
+}