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