73 lines
1.7 KiB
Python
73 lines
1.7 KiB
Python
import random
|
|
from kruskal import kruskal, _brute_mst
|
|
|
|
|
|
def test_basic():
|
|
e = [("a", "b", 1), ("b", "c", 2), ("a", "c", 4)]
|
|
w, chosen = kruskal(e, ["a", "b", "c"])
|
|
assert w == 3
|
|
assert len(chosen) == 2
|
|
|
|
|
|
def test_single_node():
|
|
w, chosen = kruskal([], ["a"])
|
|
assert w == 0 and chosen == []
|
|
|
|
|
|
def test_disconnected():
|
|
try:
|
|
kruskal([("a", "b", 1)], ["a", "b", "c"])
|
|
raise AssertionError("expected")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
def test_negative():
|
|
try:
|
|
kruskal([("a", "b", -1)], ["a", "b"])
|
|
raise AssertionError("expected")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
def test_unknown_node():
|
|
try:
|
|
kruskal([("a", "z", 1)], ["a", "b"])
|
|
raise AssertionError("expected")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
def test_random_vs_brute():
|
|
rng = random.Random(42)
|
|
for _ in range(300):
|
|
n = rng.randint(1, 6)
|
|
nodes = list(range(n))
|
|
# random edge set
|
|
possible = [(u, v) for u in range(n) for v in range(u + 1, n)]
|
|
m = rng.randint(0, len(possible))
|
|
edges = []
|
|
for u, v in rng.sample(possible, m):
|
|
edges.append((u, v, rng.randint(0, 9)))
|
|
# maybe duplicate weights fine
|
|
try:
|
|
w1, _ = kruskal(edges, nodes)
|
|
ok1 = True
|
|
except ValueError:
|
|
ok1 = False
|
|
try:
|
|
w2, _ = _brute_mst(edges, nodes)
|
|
ok2 = True
|
|
except ValueError:
|
|
ok2 = False
|
|
assert ok1 == ok2, (edges, nodes, ok1, ok2)
|
|
if ok1:
|
|
assert w1 == w2, (edges, nodes, w1, w2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for n in [x for x in globals() if x.startswith("test_")]:
|
|
globals()[n]()
|
|
print("PASS", n)
|
|
print("ALL PASS")
|