summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--file.nim100
1 files changed, 100 insertions, 0 deletions
diff --git a/file.nim b/file.nim
new file mode 100644
index 0000000..bcda055
--- /dev/null
+++ b/file.nim
@@ -0,0 +1,100 @@
+type SegTree*[T, U] = ref object
+
+ len*: int
+ arr*: seq[T]
+ shift*: seq[U]
+
+proc initSegTree*[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]:
+ 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.initSegTree(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) \ No newline at end of file