push

fenwick.nim

@@ -0,0 +1,112 @@
+type Fenwick*[T] = ref object
+  arr*: seq[T]
+
+type FenwickMod*[T] = ref object
+  arr*: seq[T]
+  m*: int64
+
+proc Fenwick*[T](f: var Fenwick[T], len: int): void =
+  new(f)
+  f.arr.newSeq(len)
+
+proc Fenwick*[T](f: var Fenwick[T], len: int, default: T): void =
+  ##ializes a fenwick tree with a constant array, f[i] = default for all i.
+  new(f)
+  f.arr.newSeq(len)
+  for i in 0..<len:
+    f.arr[i] = default
+  for i in 1..len:
+    var j = i + (i and (-i))
+    if j<=len: f.arr[j-1] += f.arr[i-1]
+
+proc Fenwick*[T](f: var Fenwick[T], default: seq[T]): void =
+  ##Converts the seq into a fenwick tree in O(default.len) time.
+  new(f)
+  f.arr = default
+  for i in 1..default.len:
+    var j = i + (i and (-i))
+    if j<=default.len: f.arr[j-1] += f.arr[i-1]
+
+proc Fenwick*[T](f: var FenwickMod[T], len: int, m: int64): void =
+  new(f)
+  f.arr.newSeq(len)
+  f.m = m
+
+proc Fenwick*[T](f: var FenwickMod[T], len: int, default: T): void =
+  ##ializes a fenwick tree with a constant array, f[i] = default for all i.
+  new(f)
+  f.arr.newSeq(len)
+  for i in 0..<len:
+    f.arr[i] = default
+  for i in 1..len:
+    var j = i + (i and (-i))
+    if j<=len: f.arr[j-1] = (f.arr[j-1] + f.arr[i-1]) mod f.m
+    
+proc Fenwick*[T](f: var FenwickMod[T], default: seq[T]): void =
+  ##Converts the seq into a fenwick tree in O(default.len) time.
+  new(f)
+  f.arr = default
+  for i in 1..len:
+    var j = i + (i and (-i))
+    if j<=len: f.arr[j-1] = (f.arr[j-1] + f.arr[i-1]) mod f.m
+
+type SomeFenwick*[T] = Fenwick[T] | FenwickMod[T]
+
+
+proc len*[T](f: SomeFenwick[T]): int = f.arr.len
+
+proc sum*[T](f: Fenwick[T], i: SomeInteger): T =
+  ##Returns f[0] + f[1] + ... + f[i]. Time O(log i).
+  var ii = i+1
+  while ii>0:
+    result += f.arr[ii-1]
+    ii -= (ii and (-ii))
+
+proc sum*[T](f: FenwickMod[T], i: SomeInteger): int64 =
+  ##Returns f[0] + f[1] + ... + f[i]. Time O(log i).
+  var ii = i+1
+  while ii>0:
+    result += f.arr[ii-1]
+    result = result mod f.m
+    ii -= (ii and (-ii))
+
+proc addTo*[T](f: Fenwick[T], i: SomeInteger, x: T): void =
+  var ii = i+1
+  while ii<=f.arr.len:
+    f.arr[ii-1] += x
+    ii += (ii and (-ii))
+
+proc addTo*[T](f: var Fenwick[T], vals: seq[T]): void =
+  ##Adds vals[i] to f[i] for each i. Time O(vals.len).
+  var vals = vals
+  for i in 1..vals.len:
+    var j = i + (i and (-i))
+    if j<=vals.len: vals[j-1] = vals[j-1] + vals[i-1]
+    f.arr[i-1] += vals[i-1]
+
+proc addTo*[T](f: FenwickMod[T], i: SomeInteger, x: T): void =
+  var ii = i+1
+  while ii<=f.arr.len:
+    f.arr[ii-1] += x
+    f.arr[ii-1] = (f.arr[ii-1] mod f.m).T
+    ii += (ii and (-ii))
+
+proc addTo*[T](f: var FenwickMod[T], vals: seq[T]): void =
+  ##Adds vals[i] to f[i] for each i. Time O(vals.len).
+  var vals = vals
+  for i in 1..vals.len:
+    var j = i + (i and (-i))
+    if j<=vals.len: vals[j-1] = (vals[j-1] + vals[i-1]) mod f.m
+    f.arr[i-1] = (f.arr[i-1] + vals[i-1]) mod f.m
+
+proc `[]`*[T](f: Fenwick[T], i: SomeInteger): T =
+  if i==0: return f.sum(0)
+  return f.sum(i) - f.sum(i-1)
+
+proc `[]`*[T](f: FenwickMod[T], i: SomeInteger): T =
+  ##Accesses a single element of the base array. O(log i)
+  if i==0: return f.sum(0)
+  return (f.sum(i) - f.sum(i-1) + f.m) mod f.m
+
+proc `[]=`*[T](f: SomeFenwick[T], i: SomeInteger, x: T): void =
+  f.addTo(i, x-f[i])