implement generics in sparse table (#27)

Author: fen-dada

st/st.mbt

@@ -1,7 +1,8 @@
 ///|
-struct ST {
+struct ST[T] {
   n : Int
-  f : Array[Array[Int]]
+  f : Array[Array[T]]
+  op: (T, T) -> T
 }
 
 ///| Build Sparse Table
@@ -17,21 +18,21 @@ struct ST {
 /// # Example
 /// ```
 /// let arr = [1,2,3,4,5,6,1,2]
-/// let _ = st(arr)
+/// let _ = from(arr, @cmp.maximum)
 /// ```
-pub fn st(data : Array[Int]) -> ST {
+pub fn[T: Default] from(data : Array[T], op: (T, T) -> T) -> ST[T] {
   let n = data.length()
   let range : Int = @math.log2(n.to_double()).to_int() + 1
-  let f = make_2d(n + 1, range)
+  let f: Array[Array[T]] = Array::makei(n+1, fn(_row: Int) {Array::make(range, T::default())} )
   for i = 0; i < n; i = i + 1 {
     f[i + 1][0] = data[i]
   }
   for j = 1; j <= range - 1; j = j + 1 {
     for i = 1; i + (1 << j) - 1 <= n; i = i + 1 {
-      f[i][j] = max(f[i][j - 1], f[i + (1 << (j - 1))][j - 1])
+      f[i][j] = op(f[i][j - 1], f[i + (1 << (j - 1))][j - 1])
     }
   }
-  ST::{ n, f }
+  ST::{ n, f, op}
 }
 
 ///| Build Binary Indexed Tree
@@ -47,16 +48,16 @@ pub fn st(data : Array[Int]) -> ST {
 /// # Example
 /// ```
 ///   let arr = [1,2,3,4,5,6,1,2]
-///   let st= @st.st(arr)
-///   inspect(st.max_elem(1, 1), content="1")
-///   inspect(st.max_elem(2, 8), content="6")
+///   let st= from(arr, @cmp.maximum)
+///   inspect(st.query(1, 1), content="1")
+///   inspect(st.query(2, 8), content="6")
 /// ```
-pub fn max_elem(self : ST, left : Int, right : Int) -> Int {
+pub fn[T] query(self : ST[T], left : Int, right : Int) -> T{
   if right < left {
     abort("Invalid range.")
   } else if left > self.n || right < 1 {
     abort("Range out of bounds.")
   }
   let k = @math.log2((right - left + 1).to_double()).to_int()
-  max(self.f[left][k], self.f[right - (1 << k) + 1][k])
+  (self.op)(self.f[left][k], self.f[right - (1 << k) + 1][k])
 }

st/st.mbti

@@ -2,11 +2,11 @@
 package "Lampese/Algorithms-Moonbit/st"
 
 // Values
-fn st(Array[Int]) -> ST
+fn[T : Default] st(Array[T], (T, T) -> T) -> ST[T]
 
 // Types and methods
-type ST
-fn ST::max_elem(Self, Int, Int) -> Int
+type ST[T]
+fn[T] ST::query(Self[T], Int, Int) -> T
 
 // Type aliases
 

st/st_test.mbt

@@ -1,7 +1,28 @@
 ///|
-test {
+test "max query" {
   let arr = [1, 2, 3, 4, 5, 6, 1, 2]
-  let st = st(arr)
-  inspect(st.max_elem(1, 1), content="1")
-  inspect(st.max_elem(2, 8), content="6")
+  let st = from(arr, @cmp.maximum)
+  inspect(st.query(1, 1), content="1")
+  inspect(st.query(2, 8), content="6")
+}
+///|
+test "min query" {
+  let arr = [1, 2, 3, 4, 5, 6, 1, 2]
+  let st = from(arr, @cmp.minimum)
+  inspect(st.query(1, 1), content="1")
+  inspect(st.query(2, 8), content="1")
+}
+///|
+test "gcd query" {
+  fn gcd(a: Int, b: Int) -> Int {
+      if b==0 {
+          a
+      } else {
+        gcd(b, a%b)
+      }
+    }
+  let arr = [4, 8, 16, 24, 48, 3]
+  let st = from(arr, gcd)
+  inspect(st.query(1, 2), content="4")
+  inspect(st.query(3, 5), content="8")
 }