add some algorithm into `array2d`

include/pocketpy/tuplelist.h

@@ -22,6 +22,7 @@ struct Tuple {
 
     Tuple(PyObject*, PyObject*);
     Tuple(PyObject*, PyObject*, PyObject*);
+    Tuple(PyObject*, PyObject*, PyObject*, PyObject*);
 
     bool is_inlined() const { return _args == _inlined; }
     PyObject*& operator[](int i){ return _args[i]; }

include/typings/array2d.pyi

@@ -1,112 +1,54 @@
-from typing import Callable, Any, Generic, TypeVar
+from typing import Callable, Any, Generic, TypeVar, Literal, overload
 
 T = TypeVar('T')
 
-class array2d(Generic[T]):
-    data: list[T]       # not available in native module
+Neighborhood = Literal['moore', 'von_neumann']
 
-    def __init__(self, n_cols: int, n_rows: int, default=None):
-        self.n_cols = n_cols
-        self.n_rows = n_rows
-        if callable(default):
-            self.data = [default() for _ in range(n_cols * n_rows)]
-        else:
-            self.data = [default] * n_cols * n_rows
-    
+class array2d(Generic[T]):
+    def __init__(self, n_cols: int, n_rows: int, default=None): ...
     @property
-    def width(self) -> int:
-        return self.n_cols
-    
+    def width(self) -> int: ...
     @property
-    def height(self) -> int:
-        return self.n_rows
-    
+    def height(self) -> int: ...
     @property
-    def numel(self) -> int:
-        return self.n_cols * self.n_rows
-
-    def is_valid(self, col: int, row: int) -> bool:
-        return 0 <= col < self.n_cols and 0 <= row < self.n_rows
+    def numel(self) -> int: ...
 
-    def get(self, col: int, row: int, default=None):
-        if not self.is_valid(col, row):
-            return default
-        return self.data[row * self.n_cols + col]
+    def is_valid(self, col: int, row: int) -> bool: ...
 
-    def __getitem__(self, index: tuple[int, int]):
-        col, row = index
-        if not self.is_valid(col, row):
-            raise IndexError(f'({col}, {row}) is not a valid index for {self!r}')
-        return self.data[row * self.n_cols + col]
+    def get(self, col: int, row: int, default=None): ...
 
-    def __setitem__(self, index: tuple[int, int], value: T):
-        col, row = index
-        if not self.is_valid(col, row):
-            raise IndexError(f'({col}, {row}) is not a valid index for {self!r}')
-        self.data[row * self.n_cols + col] = value
+    @overload
+    def __getitem__(self, index: tuple[int, int]): ...
+    @overload
+    def __getitem__(self, index: tuple[slice, slice]) -> 'array2d[T]': ...
+    @overload
+    def __setitem__(self, index: tuple[int, int], value: T): ...
+    @overload
+    def __setitem__(self, index: tuple[slice, slice], value: 'array2d[T]'): ...
 
-    def __iter__(self):
-        for row in range(self.n_rows):
-            yield [self[col, row] for col in range(self.n_cols)]
-    
-    def __len__(self):
-        return self.n_rows
-    
-    def __eq__(self, other: 'array2d') -> bool:
-        if not isinstance(other, array2d):
-            return NotImplemented
-        for i in range(self.numel):
-            if self.data[i] != other.data[i]:
-                return False
-        return True
-    
-    def __ne__(self, other: 'array2d') -> bool:
-        return not self.__eq__(other)
+    def __len__(self) -> int: ...
+    def __eq__(self, other: 'array2d') -> bool: ...
+    def __ne__(self, other: 'array2d') -> bool: ...
+    def __repr__(self): ...
 
-    def __repr__(self):
-        return f'array2d({self.n_cols}, {self.n_rows})'
+    def tolist(self) -> list[list[T]]: ...
 
-    def map(self, f: Callable[[T], Any]) -> 'array2d':
-        new_a: array2d = array2d(self.n_cols, self.n_rows)
-        for i in range(self.n_cols * self.n_rows):
-            new_a.data[i] = f(self.data[i])
-        return new_a
-    
-    def copy(self) -> 'array2d[T]':
-        new_a: array2d[T] = array2d(self.n_cols, self.n_rows)
-        new_a.data = self.data.copy()
-        return new_a
+    def map(self, f: Callable[[T], Any]) -> 'array2d': ...
+    def copy(self) -> 'array2d[T]': ...
 
-    def fill_(self, value: T) -> None:
-        for i in range(self.numel):
-            self.data[i] = value
+    def fill_(self, value: T) -> None: ...
+    def apply_(self, f: Callable[[T], T]) -> None: ...
+    def copy_(self, other: 'array2d[T] | list[T]') -> None: ...
 
-    def apply_(self, f: Callable[[T], T]) -> None:
-        for i in range(self.numel):
-            self.data[i] = f(self.data[i])
+    # algorithms
+    def count_neighbors(self, value: T, neighborhood: Neighborhood = 'moore') -> 'array2d[int]':
+        """Counts the number of neighbors with the given value for each cell."""
 
-    def copy_(self, other: 'array2d[T] | list[T]') -> None:
-        if isinstance(other, list):
-            assert len(other) == self.numel
-            self.data = other.copy()
-            return
-        self.n_cols = other.n_cols
-        self.n_rows = other.n_rows
-        self.data = other.data.copy()
+    def count(self, value: T) -> int:
+        """Counts the number of cells with the given value."""
 
-    # for cellular automata
-    def count_neighbors(self, value) -> 'array2d[int]':
-        new_a = array2d(self.n_cols, self.n_rows)
-        for j in range(self.n_rows):
-            for i in range(self.n_cols):
-                count = 0
-                count += int(self.is_valid(i-1, j-1) and self[i-1, j-1] == value)
-                count += int(self.is_valid(i, j-1) and self[i, j-1] == value)
-                count += int(self.is_valid(i+1, j-1) and self[i+1, j-1] == value)
-                count += int(self.is_valid(i-1, j) and self[i-1, j] == value)
-                count += int(self.is_valid(i+1, j) and self[i+1, j] == value)
-                count += int(self.is_valid(i-1, j+1) and self[i-1, j+1] == value)
-                count += int(self.is_valid(i, j+1) and self[i, j+1] == value)
-                count += int(self.is_valid(i+1, j+1) and self[i+1, j+1] == value)
-                new_a[i, j] = count
-        return new_a
+    def find_bounding_rect(self, value: T) -> tuple[int, int, int, int] | None:
+        """Finds the bounding rectangle of the given value.
+        
+        Returns a tuple `(x, y, width, height)` or `None` if the value is not found.
+        """

src/array2d.cpp

@@ -82,37 +82,82 @@ struct Array2d{
             return self._get(col, row);
         });
 
+        #define HANDLE_SLICE()                              \
+                int start_col, stop_col, step_col;          \
+                int start_row, stop_row, step_row;          \
+                vm->parse_int_slice(PK_OBJ_GET(Slice, xy[0]), self.n_cols, start_col, stop_col, step_col);  \
+                vm->parse_int_slice(PK_OBJ_GET(Slice, xy[1]), self.n_rows, start_row, stop_row, step_row);  \
+                if(step_col != 1 || step_row != 1) vm->ValueError("slice step must be 1");  \
+                int slice_width = stop_col - start_col; \
+                int slice_height = stop_row - start_row;    \
+                if(slice_width <= 0 || slice_height <= 0) vm->ValueError("slice width and height must be positive");
+
         vm->bind__getitem__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0, PyObject* _1){
             Array2d& self = PK_OBJ_GET(Array2d, _0);
             const Tuple& xy = CAST(Tuple&, _1);
-            int col = CAST(int, xy[0]);
-            int row = CAST(int, xy[1]);
-            if(!self.is_valid(col, row)){
-                vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
+            i64 col, row;
+            if(try_cast_int(xy[0], &col) && try_cast_int(xy[1], &row)){
+                if(!self.is_valid(col, row)){
+                    vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
+                }
+                return self._get(col, row);
             }
-            return self._get(col, row);
+
+            if(is_non_tagged_type(xy[0], VM::tp_slice) && is_non_tagged_type(xy[1], VM::tp_slice)){
+                HANDLE_SLICE();
+                PyObject* new_array_obj = vm->heap.gcnew<Array2d>(Array2d::_type(vm));
+                Array2d& new_array = PK_OBJ_GET(Array2d, new_array_obj);
+                new_array.init(stop_col - start_col, stop_row - start_row);
+                for(int j = start_row; j < stop_row; j++){
+                    for(int i = start_col; i < stop_col; i++){
+                        new_array._set(i - start_col, j - start_row, self._get(i, j));
+                    }
+                }
+                return new_array_obj;
+            }
+            vm->TypeError("expected `tuple[int, int]` or `tuple[slice, slice]` as index");
+            PK_UNREACHABLE();
         });
 
         vm->bind__setitem__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0, PyObject* _1, PyObject* _2){
             Array2d& self = PK_OBJ_GET(Array2d, _0);
             const Tuple& xy = CAST(Tuple&, _1);
-            int col = CAST(int, xy[0]);
-            int row = CAST(int, xy[1]);
-            if(!self.is_valid(col, row)){
-                vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
+            i64 col, row;
+            if(try_cast_int(xy[0], &col) && try_cast_int(xy[1], &row)){
+                if(!self.is_valid(col, row)){
+                    vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
+                }
+                self._set(col, row, _2);
+                return;
             }
-            self._set(col, row, _2);
+
+            if(is_non_tagged_type(xy[0], VM::tp_slice) && is_non_tagged_type(xy[1], VM::tp_slice)){
+                HANDLE_SLICE();
+                Array2d& other = CAST(Array2d&, _2);        // _2 must be an array2d
+                if(slice_width != other.n_cols || slice_height != other.n_rows){
+                    vm->ValueError("array2d size does not match the slice size");
+                }
+                for(int j = 0; j < slice_height; j++){
+                    for(int i = 0; i < slice_width; i++){
+                        self._set(i + start_col, j + start_row, other._get(i, j));
+                    }
+                }
+                return;
+            }
+            vm->TypeError("expected `tuple[int, int]` or `tuple[slice, slice]` as index");
         });
 
-        vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
-            Array2d& self = PK_OBJ_GET(Array2d, _0);
+        #undef HANDLE_SLICE
+
+        vm->bind(type, "tolist(self)", [](VM* vm, ArgsView args){
+            Array2d& self = PK_OBJ_GET(Array2d, args[0]);
             List t(self.n_rows);
-            List row(self.n_cols);
             for(int j = 0; j < self.n_rows; j++){
+                List row(self.n_cols);
                 for(int i = 0; i < self.n_cols; i++) row[i] = self._get(i, j);
-                t[j] = VAR(row);    // copy
+                t[j] = VAR(std::move(row));
             }
-            return vm->py_iter(VAR(std::move(t)));
+            return VAR(std::move(t));
         });
 
         vm->bind__len__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
@@ -200,29 +245,75 @@ struct Array2d{
             return vm->True;
         });
 
-        // for cellular automata
-        vm->bind(type, "count_neighbors(self, value) -> array2d[int]", [](VM* vm, ArgsView args){
+        vm->bind(type, "count_neighbors(self, value, neighborhood='moore') -> array2d[int]", [](VM* vm, ArgsView args){
             Array2d& self = PK_OBJ_GET(Array2d, args[0]);
             PyObject* new_array_obj = vm->heap.gcnew<Array2d>(Array2d::_type(vm));
             Array2d& new_array = PK_OBJ_GET(Array2d, new_array_obj);
             new_array.init(self.n_cols, self.n_rows);
             PyObject* value = args[1];
-            for(int j = 0; j < new_array.n_rows; j++){
-                for(int i = 0; i < new_array.n_cols; i++){
-                    int count = 0;
-                    count += self.is_valid(i-1, j-1) && vm->py_eq(self._get(i-1, j-1), value);
-                    count += self.is_valid(i, j-1) && vm->py_eq(self._get(i, j-1), value);
-                    count += self.is_valid(i+1, j-1) && vm->py_eq(self._get(i+1, j-1), value);
-                    count += self.is_valid(i-1, j) && vm->py_eq(self._get(i-1, j), value);
-                    count += self.is_valid(i+1, j) && vm->py_eq(self._get(i+1, j), value);
-                    count += self.is_valid(i-1, j+1) && vm->py_eq(self._get(i-1, j+1), value);
-                    count += self.is_valid(i, j+1) && vm->py_eq(self._get(i, j+1), value);
-                    count += self.is_valid(i+1, j+1) && vm->py_eq(self._get(i+1, j+1), value);
-                    new_array._set(i, j, VAR(count));
+            const Str& neighborhood = CAST(Str&, args[2]);
+            if(neighborhood == "moore"){
+                for(int j = 0; j < new_array.n_rows; j++){
+                    for(int i = 0; i < new_array.n_cols; i++){
+                        int count = 0;
+                        count += self.is_valid(i-1, j-1) && vm->py_eq(self._get(i-1, j-1), value);
+                        count += self.is_valid(i, j-1) && vm->py_eq(self._get(i, j-1), value);
+                        count += self.is_valid(i+1, j-1) && vm->py_eq(self._get(i+1, j-1), value);
+                        count += self.is_valid(i-1, j) && vm->py_eq(self._get(i-1, j), value);
+                        count += self.is_valid(i+1, j) && vm->py_eq(self._get(i+1, j), value);
+                        count += self.is_valid(i-1, j+1) && vm->py_eq(self._get(i-1, j+1), value);
+                        count += self.is_valid(i, j+1) && vm->py_eq(self._get(i, j+1), value);
+                        count += self.is_valid(i+1, j+1) && vm->py_eq(self._get(i+1, j+1), value);
+                        new_array._set(i, j, VAR(count));
+                    }
                 }
+            }else if(neighborhood == "von_neumann"){
+                for(int j = 0; j < new_array.n_rows; j++){
+                    for(int i = 0; i < new_array.n_cols; i++){
+                        int count = 0;
+                        count += self.is_valid(i, j-1) && vm->py_eq(self._get(i, j-1), value);
+                        count += self.is_valid(i-1, j) && vm->py_eq(self._get(i-1, j), value);
+                        count += self.is_valid(i+1, j) && vm->py_eq(self._get(i+1, j), value);
+                        count += self.is_valid(i, j+1) && vm->py_eq(self._get(i, j+1), value);
+                        new_array._set(i, j, VAR(count));
+                    }
+                }
+            }else{
+                vm->ValueError("neighborhood must be 'moore' or 'von_neumann'");
             }
             return new_array_obj; 
         });
+
+        vm->bind(type, "count(self, value) -> int", [](VM* vm, ArgsView args){
+            Array2d& self = PK_OBJ_GET(Array2d, args[0]);
+            PyObject* value = args[1];
+            int count = 0;
+            for(int i = 0; i < self.numel; i++) count += vm->py_eq(self.data[i], value);
+            return VAR(count);
+        });
+
+        vm->bind(type, "find_bounding_rect(self, value)", [](VM* vm, ArgsView args){
+            Array2d& self = PK_OBJ_GET(Array2d, args[0]);
+            PyObject* value = args[1];
+            int left = self.n_cols;
+            int top = self.n_rows;
+            int right = 0;
+            int bottom = 0;
+            for(int j = 0; j < self.n_rows; j++){
+                for(int i = 0; i < self.n_cols; i++){
+                    if(vm->py_eq(self._get(i, j), value)){
+                        left = std::min(left, i);
+                        top = std::min(top, j);
+                        right = std::max(right, i);
+                        bottom = std::max(bottom, j);
+                    }
+                }
+            }
+            int width = right - left + 1;
+            int height = bottom - top + 1;
+            if(width <= 0 || height <= 0) return vm->None;
+            return VAR(Tuple(VAR(left), VAR(top), VAR(width), VAR(height)));
+        });
     }
 
     void _gc_mark() const{

src/tuplelist.cpp

@@ -44,6 +44,13 @@ Tuple::Tuple(PyObject* _0, PyObject* _1, PyObject* _2): Tuple(3){
     _args[2] = _2;
 }
 
+Tuple::Tuple(PyObject* _0, PyObject* _1, PyObject* _2, PyObject* _3): Tuple(4){
+    _args[0] = _0;
+    _args[1] = _1;
+    _args[2] = _2;
+    _args[3] = _3;
+}
+
 Tuple::~Tuple(){ if(!is_inlined()) pool64_dealloc(_args); }
 
 List ArgsView::to_list() const{

tests/83_array2d.py

@@ -50,7 +50,7 @@ except IndexError:
 
 # test __iter__
 a_list = [[5, 0], [0, 0], [0, 0], [0, 6]]
-assert a_list == list(a)
+assert a_list == a.tolist()
 
 # test __len__
 assert len(a) == 4
@@ -68,8 +68,8 @@ assert repr(a) == f'array2d(2, 4)'
 
 # test map
 c = a.map(lambda x: x + 1)
-assert list(c) == [[6, 1], [1, 1], [1, 1], [1, 7]]
-assert list(a) == [[5, 0], [0, 0], [0, 0], [0, 6]]
+assert c.tolist() == [[6, 1], [1, 1], [1, 1], [1, 7]]
+assert a.tolist() == [[5, 0], [0, 0], [0, 0], [0, 6]]
 assert c.width == c.n_cols == 2
 assert c.height == c.n_rows == 4
 assert c.numel == 8
@@ -109,3 +109,26 @@ assert A().get(0, 0, default=2) == 0
 a = array2d(3, 3, default=0)
 a.count_neighbors(0) == a
 
+# test slice get
+a = array2d(5, 5, default=0)
+b = array2d(3, 2, default=1)
+
+assert a[1:4, 1:4] == array2d(3, 3, default=0)
+assert a[1:4, 1:3] == array2d(3, 2, default=0)
+assert a[1:4, 1:3] != b
+a[1:4, 1:3] = b
+assert a[1:4, 1:3] == b
+"""
+0 0 0 0 0
+0 1 1 1 0
+0 1 1 1 0
+0 0 0 0 0
+0 0 0 0 0
+"""
+assert a.count(1) == 3*2
+
+assert a.find_bounding_rect(1) == (1, 1, 3, 2)
+assert a.find_bounding_rect(0) == (0, 0, 5, 5)
+assert a.find_bounding_rect(2) == None
+
+