push

segtree.nim

@@ -0,0 +1,98 @@
+type SegTree*[T, U] = ref object
+  len*: int
+  arr*: seq[T]
+  shift*: seq[U]
+
+proc SegTree*[T, U](f: var SegTree[T, U], len: int): void =
+  new(f)
+  f.len = len
+  f.arr.newSeq(4*len)
+  f.shift.newSeq(4*len)
+
+proc build[T, U](f: var SegTree[T, U], v, tl, tr: int, a: seq[T]): void =
+    if tl==tr:
+      f.arr[v] = a[tl]
+    else:
+      var tm: int = (tl + tr) div 2
+      f.build(2*v, tl, tm, a)
+      f.build(2*v+1, tm+1, tr, a)
+      f.arr[v] = f.arr[2*v] + f.arr[2*v+1]    
+
+proc build*[T, U](f: var SegTree[T, U], a: seq[T]): void =
+  f.build(1, 0, f.len - 1, a)
+  
+proc modifyRange[T, U](tl, tr: int, sum: T, x: U, additive: bool): T =
+  if additive: return sum @ ((tr-tl+1) * x)
+  else: return sum @ x
+
+proc sum[T, U, R](f: var SegTree[T, U], v, tl, tr, l, r: int, phi: proc (x: T): R, additive: bool): R =
+  if l>r:
+    return
+  if f.shift[v] - f.shift[v] != f.shift[v]: #nonzero
+    f.arr[v] = modifyRange(tl, tr, f.arr[v], f.shift[v], additive)
+    if tl!=tr:
+      f.shift[2*v] = f.shift[2*v] + f.shift[v]
+      f.shift[2*v+1] = f.shift[2*v+1] + f.shift[v]
+    f.shift[v] = f.shift[v] - f.shift[v]
+  if l<=tl and tr<=r: return f.arr[v].phi
+  elif tl!=tr:
+    var tm: int = (tl + tr) div 2
+    return f.sum(2*v, tl, tm, l, min(r, tm), phi, additive) + f.sum(2*v+1, tm+1, tr, max(l, tm+1), r, phi, additive)
+
+proc sum*[T, U, R](f: var SegTree[T, U], i, j: int, phi: proc (x: T): R, additive: bool): R =
+  return f.sum(1, 0, f.len - 1, i, j, phi, additive)
+
+proc showDetails*[T, U](f: var SegTree[T, U], v: int = 1, tl: int = 0, tr: int = f.len - 1, prefix: string = ""): void =
+  echo prefix, "(", tl, ", ", tr, "): ", f.arr[v], ", ", f.shift[v]
+  if tl != tr:
+    var tm: int = (tl + tr) div 2
+    f.showDetails(2*v, tl, tm, prefix & "  ") 
+    f.showDetails(2*v+1, tm+1, tr, prefix & "  ")
+
+
+proc applyTo[T, U](f: var SegTree[T, U], v, tl, tr, l, r: int, x: U, additive: bool): void =
+  if f.shift[v] - f.shift[v] != f.shift[v]: #nonzero
+    f.arr[v] = modifyRange(tl, tr, f.arr[v], f.shift[v], additive)
+    if tl != tr:
+      f.shift[2*v] = f.shift[2*v] + f.shift[v]
+      f.shift[2*v+1] = f.shift[2*v+1] + f.shift[v]
+    f.shift[v] = f.shift[v] - f.shift[v]
+  if tl>tr or tl>r or tr<l: return
+  if l<=tl and tr<=r:
+    f.arr[v] = modifyRange(tl, tr, f.arr[v], x, additive)
+    if tl != tr:
+      f.shift[2*v] = f.shift[2*v] + x
+      f.shift[2*v+1] = f.shift[2*v+1] + x
+  elif tl!=tr:
+    var tm: int = (tl + tr) div 2
+    f.applyTo(2*v, tl, tm, l, min(r, tm), x, additive)
+    f.applyTo(2*v+1, tm+1, tr, max(l, tm+1), r, x, additive)
+    f.arr[v] = f.arr[2*v] + f.arr[2*v+1]
+
+proc addTo*[T, U](f: var SegTree[T, U], i, j: int, x: U): void =
+
+  f.applyTo(1, 0, f.len - 1, i, j, x, true)
+
+proc multTo*[T, U](f: var SegTree[T, U], i, j: int, x: U): void =
+
+  f.applyTo(1, 0, f.len - 1, i, j, x, false)
+
+proc toBaseArray*[T, U, R](f: var SegTree[T, U], phi: proc (x: T): R, additive: bool): seq[R] =
+  result.newSeq(f.len)
+  for i in 0..<f.len:
+    result[i] = f.sum(i, i, phi, additive)
+
+when isMainModule:
+  proc `@`(x, y: int): int = x + y
+  var seg: SegTree[int, int]
+  seg.SegTree(10)
+  var a: seq[int] =  @[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+  seg.build(a)
+  seg.showDetails
+  proc phi(x: int): int = x 
+  seg.addTo(1, 8, 4)
+  seg.showDetails
+  seg.addTo(3, 9, 4)
+  seg.showDetails
+  for i in 0..9:
+    echo i, ": ", seg.sum(i, i, phi, true)