def test_dsu(self):
number_of_categories = 100 # initial number_of_categories
new_number_of_categories = 10 # will be merged into new_number_of_categories later
number_of_items = number_of_categories + 10 ** 5
max_item = number_of_items + 10 ** 8
number_of_test = 10 ** 3
items = sample(range(max_item), number_of_items)
sets = [[item] for item in items[:number_of_categories]] # make each category non-empty
for item in items[number_of_categories:]: # then populate
choice(sets).append(item)
dsu = DSUGeneral(items)
for (leader, s) in ((choice(s), s) for s in sets): # for each category pick leader
for item in s: # and merge all items in category
dsu.union(item, leader)
self.check_dsu(dsu, sets, number_of_test)
shuffle(sets)
sets, tail = sets[:new_number_of_categories], sets[new_number_of_categories:]
for (s, t) in ((choice(sets), t) for t in tail):
dsu.union(choice(s), choice(t))
s.extend(t)
self.check_dsu(dsu, sets, number_of_test)
def check_spanning(self, tree, graph):
dsu = DSUGeneral(tree._vert)
for v in tree._vert:
for u, _ in tree._vert[v]:
dsu.union(u, v)
self.assertTrue(dsu.single_set())
self.assertEqual(len(tree._edge), len(graph._vert) - 1)
def mst(self):
res = WeightedGraph()
dsu = DSUGeneral(self._vert)
edges = sorted(self._edge.items(), key=operator.itemgetter(1))
for ((v1, v2), weight) in edges:
if not dsu.same_set(v1, v2):
dsu.union(v1, v2)
res.add_direct_link(v1, v2, weight)
return res
