class TestMinimumSpanningTree(unittest.TestCase):
def setUp(self):
self.g = Graph()
self.g.add_connected_vertice(0, 0)
self.g.add_connected_vertice(3, 4)
self.g.add_connected_vertice(0, 4)
self.g.add_connected_vertice(3, 0)
self.g.add_connected_vertice(10, 40)
self.mst = [[
Vertice(0, 0),
Vertice(0, 4),
Vertice(3, 0),
Vertice(3, 4),
Vertice(10, 40),
]]
self.three_groups = [
[Vertice(0, 0), Vertice(3, 0)],
[Vertice(3, 4), Vertice(0, 4)],
[Vertice(10, 40)],
]
def test_kruskal(self):
mst = kruskal(self.g)
self.assertGroupsEqual(mst, self.mst)
def test_kruskal_groups(self):
three_groups = kruskal(self.g, 3)
self.assertGroupsEqual(three_groups, self.three_groups)
def test_prim(self):
mst = prim(self.g)
self.assertGroupsEqual(mst, self.mst)
def test_prim_groups(self):
three_groups = prim(self.g, 3)
self.assertGroupsEqual(three_groups, self.three_groups)
def assertGroupsEqual(self, first, second, msg=None):
try:
key = lambda x: len(x)
first.sort(key=key)
second.sort(key=key)
self.assertEqual(len(first), len(second))
for i, group in enumerate(first):
self.assertEqual(len(group), len(second[i]))
for item in group:
self.assertIn(item, second[i])
except AssertionError:
msg = "Groups not equal\n\n{}\n{}".format(first, second)
raise AssertionError(msg)
class TestGraph(unittest.TestCase):
def setUp(self):
self.graph = Graph()
def test_add_connected_vertice(self):
vertices = []
edges = []
self.assertEqual(self.graph.edges, edges)
self.assertEqual(self.graph.vertices, vertices)
self.graph.add_connected_vertice(0, 0)
vertices.append(Vertice(0, 0))
self.assertEqual(self.graph.edges, edges)
self.assertEqual(self.graph.vertices, vertices)
self.graph.add_connected_vertice(3, 4)
vertices.append(Vertice(3, 4))
edges.append(Edge(vertices[0], vertices[1]))
self.assertEqual(self.graph.edges, edges)
self.graph.add_connected_vertice(0, 4)
vertices.append(Vertice(0, 4))
edges.append(Edge(vertices[0], vertices[2]))
edges.append(Edge(vertices[1], vertices[2]))
self.assertEqual(self.graph.edges, edges)
self.assertEqual(self.graph.vertices, vertices)
def test_repr(self):
self.graph.add_connected_vertice(3, 4)
self.assertEqual(str(self.graph), repr(self.graph))
def test_graphviz(self):
self.graph.add_connected_vertice(0, 0)
self.graph.add_connected_vertice(3, 4)
self.graph.add_connected_vertice(0, 4)
self.graph.edges[-1].mst = True
self.graph.add_connected_vertice(3, 0)
graphviz_res = """
graph { "(0, 0)" -- "(3, 4)"[label=5.0,len=2.3219];
"(0, 0)" -- "(0, 4)"[label=4.0,len=2.0000];
"(3, 4)" -- "(0, 4)"[label=3.0,len=1.5850,color=red,penwidth=2.0];
"(0, 0)" -- "(3, 0)"[label=3.0,len=1.5850];
"(3, 4)" -- "(3, 0)"[label=4.0,len=2.0000];
"(0, 4)" -- "(3, 0)"[label=5.0,len=2.3219]; }""".strip()
print('-' * 80)
print(graphviz_res)
print('-' * 80)
print(self.graph.to_graphviz())
print('-' * 80)
assert_equal(self.graph.to_graphviz(), graphviz_res)
def test_vertice_index(self):
self.graph.add_connected_vertice(0, 0)
vertice = self.graph.vertices[-1]
self.assertEqual(self.graph.index(vertice), 0)
self.graph.add_connected_vertice(3, 4)
vertice = self.graph.vertices[-1]
self.assertEqual(self.graph.index(vertice), 1)
self.graph.add_connected_vertice(0, 4)
vertice = self.graph.vertices[-1]
self.assertEqual(self.graph.index(vertice), 2)