1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38
39. 39
40. 40

namespace Fenwick

def lsb (i : Nat) := (i.toInt64 &&& -i.toInt64).toNatClampNeg

-- https://cp-algorithms.com/data_structures/fenwick.html

def build [Add α] (A : Vector α n) := Id.run do
  let mut F := A
  for hi : i in [1:n] do
    let j := i + lsb i
    if hj : j < n then
      F := F.set j <| F[j] + F[i]
  return F

def update [Add α] (F : Vector α n) i v := Id.run do
  let mut i := i
  let mut F := F
  while hi : i < n do
    F := F.set i <| F[i] + v
    i := i + lsb i
  return F

def query [Zero α] [Add α] (F : Vector α n) i := Id.run do
  let mut i := i
  let mut ret := 0
  while hi : 0 < i ∧ i < n do
    ret := ret + F[i]
    i := i - lsb i
  return ret

def search [Sub α] [LE α] [DecidableLE α] [NeZero n] (F : Vector α n) s := Id.run do
  let mut i := 2 ^ n.log2
  let mut j := (0 : Fin n)
  let mut s := s
  while hi : 0 < i do
    if h : j + i < n then if F[j + i] ≤ s then
      j := j + (.ofNat n i)
      s := s - F[j]
    i := i / 2
  return j