Bladeren bron

add `indexed_apply_`

blueloveTH 1 jaar geleden
bovenliggende
commit
4f431f4095
3 gewijzigde bestanden met toevoegingen van 17 en 0 verwijderingen
  1. 1 0
      include/typings/array2d.pyi
  2. 10 0
      src/array2d.cpp
  3. 6 0
      tests/83_array2d.py

+ 1 - 0
include/typings/array2d.pyi

@@ -38,6 +38,7 @@ class array2d(Generic[T]):
 
     def fill_(self, value: T) -> None: ...
     def apply_(self, f: Callable[[T], T]) -> None: ...
+    def indexed_apply_(self, f: Callable[[int, int, T], T]) -> None: ...
     def copy_(self, other: 'array2d[T] | list[T]') -> None: ...
 
     # algorithms

+ 10 - 0
src/array2d.cpp

@@ -229,6 +229,16 @@ struct Array2d{
             return vm->None;
         });
 
+        vm->bind(type, "indexed_apply_(self, f)", [](VM* vm, ArgsView args){
+            Array2d& self = PK_OBJ_GET(Array2d, args[0]);
+            PyObject* f = args[1];
+            for(int i = 0; i < self.numel; i++){
+                std::div_t res = std::div(i, self.n_cols);
+                self.data[i] = vm->call(f, VAR(res.rem), VAR(res.quot), self.data[i]);
+            }
+            return vm->None;
+        });
+
         vm->bind(type, "copy_(self, other)", [](VM* vm, ArgsView args){
             Array2d& self = PK_OBJ_GET(Array2d, args[0]);
             if(is_type(args[1], VM::tp_list)){

+ 6 - 0
tests/83_array2d.py

@@ -166,3 +166,9 @@ try:
     exit(1)
 except TypeError:
     pass
+
+a = array2d(3, 4)
+a.indexed_apply_(lambda x, y, val: x+y)
+assert a[0, 0] == 0
+assert a[1, 2] == 3
+assert a[2, 0] == 2