feat: 二次元累積和ライブラリを作成

Author: hidehic0

code/main.py

@@ -1585,6 +1585,40 @@ def update(self, i: int, x: Any) -> None:
         self.blocks[block_ind] = self.op(self.blocks[block_ind], self.lis[i])
 
 
+class PrefixSum2D:
+    def __init__(self, h: int, w: int) -> None:
+        self.data = [[0] * (w + 1) for _ in [0] * (h + 1)]
+        self.builded = False
+        self.h = h
+        self.w = w
+
+    def add(self, x: int, y: int, a: int) -> None:
+        assert 0 <= x < self.h and 0 <= y < self.w
+
+        self.data[x + 1][y + 1] += a
+
+    def build(self) -> None:
+        assert not self.builded
+
+        for i in range(self.h + 1):
+            for k in range(self.w):
+                self.data[i][k + 1] += self.data[i][k]
+
+        for k in range(self.w + 1):
+            for i in range(self.h):
+                self.data[i + 1][k] += self.data[i][k]
+
+    def prod(self, ax: int, ay: int, bx: int, by: int) -> int:
+        assert 0 <= ax <= bx < self.h and 0 <= ay <= by < self.w
+
+        return (
+            self.data[bx + 1][by + 1]
+            + self.data[ax][ay]
+            - self.data[ax][by + 1]
+            - self.data[bx + 1][ay]
+        )
+
+
 from typing import Any, Callable
 
 

libs/prefix_sum_2d.py

@@ -0,0 +1,32 @@
+class PrefixSum2D:
+    def __init__(self, h: int, w: int) -> None:
+        self.data = [[0] * (w + 1) for _ in [0] * (h + 1)]
+        self.builded = False
+        self.h = h
+        self.w = w
+
+    def add(self, x: int, y: int, a: int) -> None:
+        assert 0 <= x < self.h and 0 <= y < self.w
+
+        self.data[x + 1][y + 1] += a
+
+    def build(self) -> None:
+        assert not self.builded
+
+        for i in range(self.h + 1):
+            for k in range(self.w):
+                self.data[i][k + 1] += self.data[i][k]
+
+        for k in range(self.w + 1):
+            for i in range(self.h):
+                self.data[i + 1][k] += self.data[i][k]
+
+    def prod(self, ax: int, ay: int, bx: int, by: int) -> int:
+        assert 0 <= ax <= bx < self.h and 0 <= ay <= by < self.w
+
+        return (
+            self.data[bx + 1][by + 1]
+            + self.data[ax][ay]
+            - self.data[ax][by + 1]
+            - self.data[bx + 1][ay]
+        )

merge_file.bash

@@ -7,7 +7,7 @@ echo "新しいmain.py作成完了"
 # テンプレ
 # /bin/cat python/<filename>.py >> code/main.py
 
-for file_name in "import.py" "standard_input.py" "math_func.py" "array_create.py" "binary_search.py" "modint.py" "yn_func.py" "grid.py" "coordinates_to_id.py" "dijkstra.py" "get_path.py" "dp.py" "coordinate_compression.py" "memo.py" "rerooting.py" "manacher.py" "rollinghash.py" "graph.py" "unionfind.py" "euler_tour.py" "potential_unionfind.py" "heap.py" "trie.py" "mo.py" "square-division.py" "dual_segtree.py" "dis_lib.py" "alias.py" "utils.py"; do
+for file_name in "import.py" "standard_input.py" "math_func.py" "array_create.py" "binary_search.py" "modint.py" "yn_func.py" "grid.py" "coordinates_to_id.py" "dijkstra.py" "get_path.py" "dp.py" "coordinate_compression.py" "memo.py" "rerooting.py" "manacher.py" "rollinghash.py" "graph.py" "unionfind.py" "euler_tour.py" "potential_unionfind.py" "heap.py" "trie.py" "mo.py" "square-division.py" "prefix_sum_2d.py" "dual_segtree.py" "dis_lib.py" "alias.py" "utils.py"; do
 
   cat libs/${file_name} >>code/main.py
 done

tests/aoj0560.py

@@ -0,0 +1,32 @@
+from libs.prefix_sum_2d import PrefixSum2D
+
+H, W = map(int, input().split())
+K = int(input())
+S = [input() for _ in [0] * H]
+
+JS, IS, OS = PrefixSum2D(H, W), PrefixSum2D(H, W), PrefixSum2D(H, W)
+
+for i in range(H):
+    for k in range(W):
+        match S[i][k]:
+            case "J":
+                JS.add(i, k, 1)
+
+            case "O":
+                OS.add(i, k, 1)
+
+            case "I":
+                IS.add(i, k, 1)
+
+JS.build()
+IS.build()
+OS.build()
+
+for _ in [0] * K:
+    ax, ay, bx, by = map(int, input().split())
+    ax -= 1
+    ay -= 1
+    bx -= 1
+    by -= 1
+
+    print(JS.prod(ax, ay, bx, by), OS.prod(ax, ay, bx, by), IS.prod(ax, ay, bx, by))