From 1ca69b3a355693c1756a2c46a8178c13f0c9c8a6 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Mon, 25 Nov 2024 12:42:42 +0800
Subject: [PATCH] ...

---
 include/typings/array2d.pyi |  7 +++++++
 src/modules/array2d.c       |  8 ++++++++
 tests/90_array2d.py         | 27 +++++++++++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi
index dae03220..2a7bd517 100644
--- a/include/typings/array2d.pyi
+++ b/include/typings/array2d.pyi
@@ -84,3 +84,10 @@ class array2d(Generic[T]):
 
     def convolve(self: array2d[int], kernel: array2d[int], padding: int) -> array2d[int]:
         """Convolves the array with the given kernel."""
+
+    def get_connected_components(self, value: T, neighborhood: Neighborhood) -> tuple[array2d[int], int]:
+        """Gets connected components of the grid.
+
+        Returns the `visited` array and the number of connected components,
+        where `0` means unvisited, and non-zero means the index of the connected component.
+        """
diff --git a/src/modules/array2d.c b/src/modules/array2d.c
index e1784ab1..aa07d3d6 100644
--- a/src/modules/array2d.c
+++ b/src/modules/array2d.c
@@ -764,6 +764,14 @@ void pk__add_module_array2d() {
     py_bindmethod(array2d, "get_bounding_rect", array2d_get_bounding_rect);
     py_bindmethod(array2d, "count_neighbors", array2d_count_neighbors);
     py_bindmethod(array2d, "convolve", array2d_convolve);
+
+    const char* scc =
+        "\ndef get_connected_components(self, value: T, neighborhood: Neighborhood) -> tuple[array2d[int], int]:\n    from collections import deque\n    from linalg import vec2i\n\n    DIRS = [vec2i.LEFT, vec2i.RIGHT, vec2i.UP, vec2i.DOWN]\n    assert neighborhood in ['Moore', 'von Neumann']\n\n    if neighborhood == 'Moore':\n        DIRS.extend([\n            vec2i.LEFT+vec2i.UP,\n            vec2i.RIGHT+vec2i.UP,\n            vec2i.LEFT+vec2i.DOWN,\n            vec2i.RIGHT+vec2i.DOWN\n            ])\n\n    visited = array2d[int](self.width, self.height, default=0)\n    queue = deque()\n    count = 0\n    for y in range(self.height):\n        for x in range(self.width):\n            if visited[x, y] or self[x, y] != value:\n                continue\n            count += 1\n            queue.append((x, y))\n            visited[x, y] = count\n            while queue:\n                cx, cy = queue.popleft()\n                for dx, dy in DIRS:\n                    nx, ny = cx+dx, cy+dy\n                    if self.is_valid(nx, ny) and not visited[nx, ny] and self[nx, ny] == value:\n                        queue.append((nx, ny))\n                        visited[nx, ny] = count\n    return visited, count\n\narray2d.get_connected_components = get_connected_components\ndel get_connected_components\n";
+
+    if(!py_exec(scc, "array2d.py", EXEC_MODE, mod)) {
+        py_printexc();
+        c11__abort("failed to execute array2d.py");
+    }
 }
 
 #undef INC_COUNT
diff --git a/tests/90_array2d.py b/tests/90_array2d.py
index bb3d6e30..9f4b3b7b 100644
--- a/tests/90_array2d.py
+++ b/tests/90_array2d.py
@@ -217,6 +217,33 @@ assert res[mask] == [0, 4, 5, 0, 4, 5]
 res[mask] = -1
 assert res.tolist() == [[-1, -1, 9, 9, -1], [-1, -1, 9, 9, -1]]
 
+# test get_connected_components
+a = array2d[int].fromlist([
+    [1, 1, 0, 1],
+    [0, 2, 2, 1],
+    [0, 1, 1, 1],
+    [1, 0, 0, 0],
+])
+vis, cnt = a.get_connected_components(1, 'von Neumann')
+assert vis == [
+    [1, 1, 0, 2],
+    [0, 0, 0, 2],
+    [0, 2, 2, 2],
+    [3, 0, 0, 0]
+    ]
+assert cnt == 3
+vis, cnt = a.get_connected_components(1, 'Moore')
+assert vis == [
+    [1, 1, 0, 2],
+    [0, 0, 0, 2],
+    [0, 2, 2, 2],
+    [2, 0, 0, 0]
+    ]
+assert cnt == 2
+vis, cnt = a.get_connected_components(2, 'von Neumann')
+assert cnt == 1
+vis, cnt = a.get_connected_components(0, 'Moore')
+assert cnt == 2
 
 # stackoverflow bug due to recursive mark-and-sweep
 # class Cell: