diff --git a/cyaron/graph.py b/cyaron/graph.py index c5965ed..3b02c99 100644 --- a/cyaron/graph.py +++ b/cyaron/graph.py @@ -129,6 +129,36 @@ def add_edge(self, x, y, **kwargs): if not self.directed and x != y: self.__add_edge(y, x, weight) + def to_tree(self, root = 1, *, container = list, from_weight = lambda w: None): + """to_tree(self, root, *, container) -> list[tuple[int, int]] + Return a spanning tree of the graph. + By default, return the DFS spanning tree generated from node 1 as the root. + int root = 1 -> starting from `root` to generate. + container = list -> use list to simulate a stack by default. + You can pass in a container class that implements `pop` and `append` methods. + Additionally, it needs the method `__bool__` to indicate whether the container is non-empty. + For example, pass in a heap to implement a minimum spanning tree. + from_weight = lambda w: None -> use to print the weight of edge, if it returns `None`, the function will ignore the weight. + """ + tree = [] + vis = set() + q = container() + q.append((0, root, 0)) + while q: + _, u, fa = q.pop() + if u in vis: + continue + vis.add(u) + if fa != 0: + w = from_weight(_) + tree.append((fa, u) if w is None else (fa, u, w)) + for edge in self.edges[u]: + v, w = edge.end, edge.weight + if v in vis: + continue + q.append((w, v, u)) + return tree + @staticmethod def chain(point_count, **kwargs): """chain(point_count, **kwargs) -> Graph