112 lines
3.3 KiB
Nim
112 lines
3.3 KiB
Nim
|
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])
|