113 lines
2.6 KiB
Python
113 lines
2.6 KiB
Python
import random
|
|
from dsu import DSU
|
|
|
|
|
|
def test_basic():
|
|
d = DSU(5)
|
|
assert d.components == 5
|
|
assert not d.connected(0, 1)
|
|
assert d.union(0, 1) is True
|
|
assert d.connected(0, 1)
|
|
assert d.components == 4
|
|
assert d.union(0, 1) is False
|
|
assert d.components == 4
|
|
|
|
|
|
def test_size():
|
|
d = DSU(6)
|
|
d.union(0, 1); d.union(1, 2)
|
|
assert d.size(0) == 3 and d.size(2) == 3
|
|
assert d.size(3) == 1
|
|
d.union(3, 4)
|
|
assert d.size(3) == 2
|
|
|
|
|
|
def test_transitive():
|
|
d = DSU(4)
|
|
d.union(0, 1); d.union(2, 3); d.union(1, 2)
|
|
assert d.components == 1
|
|
for i in range(4):
|
|
for j in range(4):
|
|
assert d.connected(i, j)
|
|
|
|
|
|
def test_empty():
|
|
d = DSU(0)
|
|
assert d.components == 0
|
|
assert len(d) == 0
|
|
assert d.sets() == {}
|
|
|
|
|
|
def test_reset():
|
|
d = DSU(4)
|
|
d.union(0, 1); d.union(2, 3)
|
|
d.reset()
|
|
assert d.components == 4
|
|
assert not d.connected(0, 1)
|
|
|
|
|
|
def test_sets():
|
|
d = DSU(5)
|
|
d.union(0, 2); d.union(1, 3)
|
|
s = d.sets()
|
|
assert sorted(s.values(), key=lambda v: v[0]) == [[0, 2], [1, 3], [4]]
|
|
|
|
|
|
def test_errors():
|
|
d = DSU(3)
|
|
try:
|
|
d.find(5)
|
|
except IndexError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected IndexError")
|
|
try:
|
|
DSU(-1)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected ValueError")
|
|
try:
|
|
DSU(True)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected ValueError for bool")
|
|
|
|
|
|
def test_random_vs_reference():
|
|
rng = random.Random(42)
|
|
n = 30
|
|
d = DSU(n)
|
|
parent = list(range(n))
|
|
def ref_find(x):
|
|
while parent[x] != x:
|
|
parent[x] = parent[parent[x]]
|
|
x = parent[x]
|
|
return x
|
|
for _ in range(2000):
|
|
a = rng.randrange(n); b = rng.randrange(n)
|
|
# union in both
|
|
ra, rb = ref_find(a), ref_find(b)
|
|
expected_merge = ra != rb
|
|
if expected_merge:
|
|
parent[ra] = rb
|
|
got = d.union(a, b)
|
|
assert got == expected_merge, (a, b, got, expected_merge)
|
|
# connectivity check sample
|
|
for _ in range(5):
|
|
x = rng.randrange(n); y = rng.randrange(n)
|
|
assert d.connected(x, y) == (ref_find(x) == ref_find(y))
|
|
assert d.components == len({ref_find(i) for i in range(n)})
|
|
|
|
|
|
def test_deep_chain_no_recursion():
|
|
# union in a line to force deep chains, then find must not recurse
|
|
n = 100000
|
|
d = DSU(n)
|
|
for i in range(n - 1):
|
|
d.union(i, i + 1)
|
|
assert d.components == 1
|
|
assert d.find(0) == d.find(n - 1)
|
|
assert d.size(0) == n
|