139 lines
3.3 KiB
Python
139 lines
3.3 KiB
Python
import random
|
|
from fenwick import Fenwick
|
|
|
|
|
|
def test_basic_update_prefix():
|
|
ft = Fenwick(5)
|
|
for i in range(1, 6):
|
|
ft.update(i, i)
|
|
assert ft.prefix(0) == 0
|
|
assert ft.prefix(1) == 1
|
|
assert ft.prefix(3) == 6
|
|
assert ft.prefix(5) == 15
|
|
assert ft.total() == 15
|
|
|
|
|
|
def test_range_sum():
|
|
ft = Fenwick(6)
|
|
arr = [3, 1, 4, 1, 5, 9]
|
|
for i, v in enumerate(arr, start=1):
|
|
ft.update(i, v)
|
|
assert ft.range_sum(2, 4) == 1 + 4 + 1
|
|
assert ft.range_sum(1, 6) == sum(arr)
|
|
assert ft.range_sum(3, 3) == 4
|
|
|
|
|
|
def test_from_array():
|
|
arr = [2, 5, 1, 3, 7]
|
|
ft = Fenwick.from_array(arr)
|
|
for i in range(1, len(arr) + 1):
|
|
assert ft.prefix(i) == sum(arr[:i])
|
|
# update works on built tree
|
|
ft.update(3, 10)
|
|
assert ft.prefix(3) == 2 + 5 + 1 + 10
|
|
|
|
|
|
def test_lower_bound():
|
|
ft = Fenwick(5)
|
|
for i in range(1, 6):
|
|
ft.update(i, 1) # all ones
|
|
# prefix: 1,2,3,4,5
|
|
assert ft.lower_bound(1) == 1
|
|
assert ft.lower_bound(2) == 2
|
|
assert ft.lower_bound(5) == 5
|
|
assert ft.lower_bound(6) == 6 # none
|
|
|
|
|
|
def test_lower_bound_weighted():
|
|
ft = Fenwick(4)
|
|
ft.update(1, 5)
|
|
ft.update(2, 2)
|
|
ft.update(3, 7)
|
|
ft.update(4, 1)
|
|
# prefixes: 5,7,14,15
|
|
assert ft.lower_bound(5) == 1
|
|
assert ft.lower_bound(6) == 2
|
|
assert ft.lower_bound(7) == 2
|
|
assert ft.lower_bound(8) == 3
|
|
assert ft.lower_bound(15) == 4
|
|
assert ft.lower_bound(16) == 5
|
|
|
|
|
|
def test_random_vs_brute():
|
|
n = 50
|
|
arr = [0] * n
|
|
ft = Fenwick(n)
|
|
rng = random.Random(42)
|
|
for _ in range(2000):
|
|
op = rng.random()
|
|
if op < 0.5:
|
|
i = rng.randrange(1, n + 1)
|
|
d = rng.randrange(-5, 6)
|
|
arr[i - 1] += d
|
|
ft.update(i, d)
|
|
else:
|
|
# check a random prefix
|
|
i = rng.randrange(0, n + 1)
|
|
assert ft.prefix(i) == sum(arr[:i]), (i, ft.prefix(i), sum(arr[:i]))
|
|
# lower_bound against non-negative only: rebuild with non-negative
|
|
arr2 = [rng.randrange(0, 6) for _ in range(n)]
|
|
ft2 = Fenwick.from_array(arr2)
|
|
# verify lower_bound for every target 1..total
|
|
total = sum(arr2)
|
|
for t in range(1, total + 1):
|
|
# brute
|
|
s = 0
|
|
brute = None
|
|
for idx, v in enumerate(arr2, start=1):
|
|
s += v
|
|
if s >= t:
|
|
brute = idx
|
|
break
|
|
if brute is None:
|
|
brute = n + 1
|
|
assert ft2.lower_bound(t) == brute, (t, ft2.lower_bound(t), brute)
|
|
|
|
|
|
def test_bounds():
|
|
ft = Fenwick(3)
|
|
try:
|
|
ft.update(0, 1)
|
|
except IndexError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected IndexError")
|
|
try:
|
|
ft.update(4, 1)
|
|
except IndexError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected IndexError")
|
|
try:
|
|
ft.prefix(4)
|
|
except IndexError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected IndexError")
|
|
try:
|
|
ft.range_sum(2, 1)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected ValueError")
|
|
|
|
|
|
def test_empty():
|
|
ft = Fenwick(0)
|
|
assert ft.total() == 0
|
|
assert ft.prefix(0) == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# run all test_ functions
|
|
import sys
|
|
g = globals()
|
|
for name in sorted(g):
|
|
if name.startswith("test_"):
|
|
g[name]()
|
|
print("ALL TESTS PASSED")
|