wehub-resource-sync
686 lines
22 KiB
Python
686 lines
22 KiB
Python
from heapq import heappop
|
|
from heapq import heappush
|
|
from itertools import count
|
|
from math import isnan
|
|
from operator import itemgetter
|
|
|
|
from easygraph.utils.decorators import *
|
|
|
|
|
|
__all__ = [
|
|
"minimum_spanning_edges",
|
|
"maximum_spanning_edges",
|
|
"minimum_spanning_tree",
|
|
"maximum_spanning_tree",
|
|
]
|
|
|
|
@hybrid("cpp_boruvka_mst_edges")
|
|
def boruvka_mst_edges(G, minimum=True, weight="weight", data=True, ignore_nan=False):
|
|
"""Iterate over edges of a Borůvka's algorithm min/max spanning tree.
|
|
|
|
Parameters
|
|
----------
|
|
G : EasyGraph Graph
|
|
The edges of `G` must have distinct weights,
|
|
otherwise the edges may not form a tree.
|
|
|
|
minimum : bool (default: True)
|
|
Find the minimum (True) or maximum (False) spanning tree.
|
|
|
|
weight : string (default: 'weight')
|
|
The name of the edge attribute holding the edge weights.
|
|
|
|
data : bool (default: True)
|
|
Flag for whether to yield edge attribute dicts.
|
|
If True, yield edges `(u, v, d)`, where `d` is the attribute dict.
|
|
If False, yield edges `(u, v)`.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
"""
|
|
# Initialize a forest, assuming initially that it is the discrete
|
|
# partition of the nodes of the graph.
|
|
forest = UnionFind(G)
|
|
|
|
def best_edge(component):
|
|
"""Returns the optimum (minimum or maximum) edge on the edge
|
|
boundary of the given set of nodes.
|
|
|
|
A return value of ``None`` indicates an empty boundary.
|
|
|
|
"""
|
|
sign = 1 if minimum else -1
|
|
minwt = float("inf")
|
|
boundary = None
|
|
for e in edge_boundary(G, component, data=True):
|
|
wt = e[-1].get(weight, 1) * sign
|
|
if isnan(wt):
|
|
if ignore_nan:
|
|
continue
|
|
msg = f"NaN found as an edge weight. Edge {e}"
|
|
raise ValueError(msg)
|
|
if wt < minwt:
|
|
minwt = wt
|
|
boundary = e
|
|
return boundary
|
|
|
|
# Determine the optimum edge in the edge boundary of each component
|
|
# in the forest.
|
|
best_edges = (best_edge(component) for component in forest.to_sets())
|
|
best_edges = [edge for edge in best_edges if edge is not None]
|
|
# If each entry was ``None``, that means the graph was disconnected,
|
|
# so we are done generating the forest.
|
|
while best_edges:
|
|
# Determine the optimum edge in the edge boundary of each
|
|
# component in the forest.
|
|
#
|
|
# This must be a sequence, not an iterator. In this list, the
|
|
# same edge may appear twice, in different orientations (but
|
|
# that's okay, since a union operation will be called on the
|
|
# endpoints the first time it is seen, but not the second time).
|
|
#
|
|
# Any ``None`` indicates that the edge boundary for that
|
|
# component was empty, so that part of the forest has been
|
|
# completed.
|
|
#
|
|
# TODO This can be parallelized, both in the outer loop over
|
|
# each component in the forest and in the computation of the
|
|
# minimum. (Same goes for the identical lines outside the loop.)
|
|
best_edges = (best_edge(component) for component in forest.to_sets())
|
|
best_edges = [edge for edge in best_edges if edge is not None]
|
|
# Join trees in the forest using the best edges, and yield that
|
|
# edge, since it is part of the spanning tree.
|
|
#
|
|
# TODO This loop can be parallelized, to an extent (the union
|
|
# operation must be atomic).
|
|
for u, v, d in best_edges:
|
|
if forest[u] != forest[v]:
|
|
if data:
|
|
yield u, v, d
|
|
else:
|
|
yield u, v
|
|
forest.union(u, v)
|
|
|
|
|
|
@hybrid("cpp_kruskal_mst_edges")
|
|
def kruskal_mst_edges(G, minimum=True, weight="weight", data=True, ignore_nan=False):
|
|
"""Iterate over edges of a Kruskal's algorithm min/max spanning tree.
|
|
|
|
Parameters
|
|
----------
|
|
G : EasyGraph Graph
|
|
The graph holding the tree of interest.
|
|
|
|
minimum : bool (default: True)
|
|
Find the minimum (True) or maximum (False) spanning tree.
|
|
|
|
weight : string (default: 'weight')
|
|
The name of the edge attribute holding the edge weights.
|
|
|
|
data : bool (default: True)
|
|
Flag for whether to yield edge attribute dicts.
|
|
If True, yield edges `(u, v, d)`, where `d` is the attribute dict.
|
|
If False, yield edges `(u, v)`.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
"""
|
|
subtrees = UnionFind()
|
|
edges = []
|
|
for u, v, t in G.edges:
|
|
edges.append((u, v, t))
|
|
|
|
def filter_nan_edges(edges=edges, weight=weight):
|
|
sign = 1 if minimum else -1
|
|
for u, v, d in edges:
|
|
wt = d.get(weight, 1) * sign
|
|
if isnan(wt):
|
|
if ignore_nan:
|
|
continue
|
|
msg = f"NaN found as an edge weight. Edge {(u, v, d)}"
|
|
raise ValueError(msg)
|
|
yield wt, u, v, d
|
|
|
|
edges = sorted(filter_nan_edges(), key=itemgetter(0))
|
|
for wt, u, v, d in edges:
|
|
if subtrees[u] != subtrees[v]:
|
|
if data:
|
|
yield (u, v, d)
|
|
else:
|
|
yield (u, v)
|
|
subtrees.union(u, v)
|
|
|
|
|
|
@hybrid("cpp_prim_mst_edges")
|
|
def prim_mst_edges(G, minimum=True, weight="weight", data=True, ignore_nan=False):
|
|
"""Iterate over edges of Prim's algorithm min/max spanning tree.
|
|
|
|
Parameters
|
|
----------
|
|
G : EasyGraph Graph
|
|
The graph holding the tree of interest.
|
|
|
|
minimum : bool (default: True)
|
|
Find the minimum (True) or maximum (False) spanning tree.
|
|
|
|
weight : string (default: 'weight')
|
|
The name of the edge attribute holding the edge weights.
|
|
|
|
data : bool (default: True)
|
|
Flag for whether to yield edge attribute dicts.
|
|
If True, yield edges `(u, v, d)`, where `d` is the attribute dict.
|
|
If False, yield edges `(u, v)`.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
"""
|
|
push = heappush
|
|
pop = heappop
|
|
|
|
nodes = set(G)
|
|
c = count()
|
|
|
|
sign = 1 if minimum else -1
|
|
|
|
while nodes:
|
|
u = nodes.pop()
|
|
frontier = []
|
|
visited = {u}
|
|
for v, d in G.adj[u].items():
|
|
wt = d.get(weight, 1) * sign
|
|
if isnan(wt):
|
|
if ignore_nan:
|
|
continue
|
|
msg = f"NaN found as an edge weight. Edge {(u, v, d)}"
|
|
raise ValueError(msg)
|
|
push(frontier, (wt, next(c), u, v, d))
|
|
while frontier:
|
|
W, _, u, v, d = pop(frontier)
|
|
if v in visited or v not in nodes:
|
|
continue
|
|
if data:
|
|
yield u, v, d
|
|
else:
|
|
yield u, v
|
|
# update frontier
|
|
visited.add(v)
|
|
nodes.discard(v)
|
|
for w, d2 in G.adj[v].items():
|
|
if w in visited:
|
|
continue
|
|
new_weight = d2.get(weight, 1) * sign
|
|
push(frontier, (new_weight, next(c), v, w, d2))
|
|
|
|
|
|
ALGORITHMS = {
|
|
"boruvka": boruvka_mst_edges,
|
|
"borůvka": boruvka_mst_edges,
|
|
"kruskal": kruskal_mst_edges,
|
|
"prim": prim_mst_edges,
|
|
}
|
|
|
|
|
|
@not_implemented_for("multigraph")
|
|
@only_implemented_for_UnDirected_graph
|
|
def minimum_spanning_edges(
|
|
G, algorithm="kruskal", weight="weight", data=True, ignore_nan=False
|
|
):
|
|
"""Generate edges in a minimum spanning forest of an undirected
|
|
weighted graph.
|
|
|
|
A minimum spanning tree is a subgraph of the graph (a tree)
|
|
with the minimum sum of edge weights. A spanning forest is a
|
|
union of the spanning trees for each connected component of the graph.
|
|
|
|
Parameters
|
|
----------
|
|
G : undirected Graph
|
|
An undirected graph. If `G` is connected, then the algorithm finds a
|
|
spanning tree. Otherwise, a spanning forest is found.
|
|
|
|
algorithm : string
|
|
The algorithm to use when finding a minimum spanning tree. Valid
|
|
choices are 'kruskal', 'prim', or 'boruvka'. The default is 'kruskal'.
|
|
|
|
weight : string
|
|
Edge data key to use for weight (default 'weight').
|
|
|
|
data : bool, optional
|
|
If True yield the edge data along with the edge.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
Returns
|
|
-------
|
|
edges : iterator
|
|
An iterator over edges in a maximum spanning tree of `G`.
|
|
Edges connecting nodes `u` and `v` are represented as tuples:
|
|
`(u, v, k, d)` or `(u, v, k)` or `(u, v, d)` or `(u, v)`
|
|
|
|
Examples
|
|
--------
|
|
>>> from easygraph.functions.basic import mst
|
|
|
|
Find minimum spanning edges by Kruskal's algorithm
|
|
|
|
>>> G.add_edge(0, 3, weight=2)
|
|
>>> mst = mst.minimum_spanning_edges(G, algorithm="kruskal", data=False)
|
|
>>> edgelist = list(mst)
|
|
>>> sorted(sorted(e) for e in edgelist)
|
|
[[0, 1], [1, 2], [2, 3]]
|
|
|
|
Find minimum spanning edges by Prim's algorithm
|
|
|
|
>>> G.add_edge(0, 3, weight=2)
|
|
>>> mst = mst.minimum_spanning_edges(G, algorithm="prim", data=False)
|
|
>>> edgelist = list(mst)
|
|
>>> sorted(sorted(e) for e in edgelist)
|
|
[[0, 1], [1, 2], [2, 3]]
|
|
|
|
Notes
|
|
-----
|
|
For Borůvka's algorithm, each edge must have a weight attribute, and
|
|
each edge weight must be distinct.
|
|
|
|
For the other algorithms, if the graph edges do not have a weight
|
|
attribute a default weight of 1 will be used.
|
|
|
|
Modified code from David Eppstein, April 2006
|
|
http://www.ics.uci.edu/~eppstein/PADS/
|
|
|
|
"""
|
|
try:
|
|
algo = ALGORITHMS[algorithm]
|
|
except KeyError as e:
|
|
msg = f"{algorithm} is not a valid choice for an algorithm."
|
|
raise ValueError(msg) from e
|
|
|
|
return algo(G, minimum=True, weight=weight, data=data, ignore_nan=ignore_nan)
|
|
|
|
|
|
@not_implemented_for("multigraph")
|
|
@only_implemented_for_UnDirected_graph
|
|
def maximum_spanning_edges(
|
|
G, algorithm="kruskal", weight="weight", data=True, ignore_nan=False
|
|
):
|
|
"""Generate edges in a maximum spanning forest of an undirected
|
|
weighted graph.
|
|
|
|
A maximum spanning tree is a subgraph of the graph (a tree)
|
|
with the maximum possible sum of edge weights. A spanning forest is a
|
|
union of the spanning trees for each connected component of the graph.
|
|
|
|
Parameters
|
|
----------
|
|
G : undirected Graph
|
|
An undirected graph. If `G` is connected, then the algorithm finds a
|
|
spanning tree. Otherwise, a spanning forest is found.
|
|
|
|
algorithm : string
|
|
The algorithm to use when finding a maximum spanning tree. Valid
|
|
choices are 'kruskal', 'prim', or 'boruvka'. The default is 'kruskal'.
|
|
|
|
weight : string
|
|
Edge data key to use for weight (default 'weight').
|
|
|
|
data : bool, optional
|
|
If True yield the edge data along with the edge.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
Returns
|
|
-------
|
|
edges : iterator
|
|
An iterator over edges in a maximum spanning tree of `G`.
|
|
Edges connecting nodes `u` and `v` are represented as tuples:
|
|
`(u, v, k, d)` or `(u, v, k)` or `(u, v, d)` or `(u, v)`
|
|
|
|
Examples
|
|
--------
|
|
>>> from easygraph.functions.path import mst
|
|
|
|
Find maximum spanning edges by Kruskal's algorithm
|
|
|
|
>>> G.add_edge(0, 3, weight=2)
|
|
>>> mst = mst.maximum_spanning_edges(G, algorithm="kruskal", data=False)
|
|
>>> edgelist = list(mst)
|
|
>>> sorted(sorted(e) for e in edgelist)
|
|
[[0, 1], [0, 3], [1, 2]]
|
|
|
|
Find maximum spanning edges by Prim's algorithm
|
|
|
|
>>> G.add_edge(0, 3, weight=2) # assign weight 2 to edge 0-3
|
|
>>> mst = mst.maximum_spanning_edges(G, algorithm="prim", data=False)
|
|
>>> edgelist = list(mst)
|
|
>>> sorted(sorted(e) for e in edgelist)
|
|
[[0, 1], [0, 3], [2, 3]]
|
|
|
|
Notes
|
|
-----
|
|
For Borůvka's algorithm, each edge must have a weight attribute, and
|
|
each edge weight must be distinct.
|
|
|
|
For the other algorithms, if the graph edges do not have a weight
|
|
attribute a default weight of 1 will be used.
|
|
|
|
Modified code from David Eppstein, April 2006
|
|
http://www.ics.uci.edu/~eppstein/PADS/
|
|
"""
|
|
try:
|
|
algo = ALGORITHMS[algorithm]
|
|
except KeyError as e:
|
|
msg = f"{algorithm} is not a valid choice for an algorithm."
|
|
raise ValueError(msg) from e
|
|
|
|
return algo(G, minimum=False, weight=weight, data=data, ignore_nan=ignore_nan)
|
|
|
|
|
|
@not_implemented_for("multigraph")
|
|
def minimum_spanning_tree(G, weight="weight", algorithm="kruskal", ignore_nan=False):
|
|
"""Returns a minimum spanning tree or forest on an undirected graph `G`.
|
|
|
|
Parameters
|
|
----------
|
|
G : undirected graph
|
|
An undirected graph. If `G` is connected, then the algorithm finds a
|
|
spanning tree. Otherwise, a spanning forest is found.
|
|
|
|
weight : str
|
|
Data key to use for edge weights.
|
|
|
|
algorithm : string
|
|
The algorithm to use when finding a minimum spanning tree. Valid
|
|
choices are 'kruskal', 'prim', or 'boruvka'. The default is
|
|
'kruskal'.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
Returns
|
|
-------
|
|
G : EasyGraph Graph
|
|
A minimum spanning tree or forest.
|
|
|
|
Examples
|
|
--------
|
|
>>> G.add_edge(0, 3, weight=2)
|
|
>>> T = eg.minimum_spanning_tree(G)
|
|
>>> sorted(T.edges(data=True))
|
|
[(0, 1, {}), (1, 2, {}), (2, 3, {})]
|
|
|
|
|
|
Notes
|
|
-----
|
|
For Borůvka's algorithm, each edge must have a weight attribute, and
|
|
each edge weight must be distinct.
|
|
|
|
For the other algorithms, if the graph edges do not have a weight
|
|
attribute a default weight of 1 will be used.
|
|
|
|
Isolated nodes with self-loops are in the tree as edgeless isolated nodes.
|
|
|
|
"""
|
|
edges = list(
|
|
minimum_spanning_edges(G, algorithm, weight, data=True, ignore_nan=ignore_nan)
|
|
)
|
|
T = G.__class__() # Same graph class as G
|
|
for i in G.nodes:
|
|
T.add_node(i)
|
|
for i in edges:
|
|
(u, v, t) = i
|
|
T.add_edge(u, v, **t)
|
|
return T
|
|
|
|
|
|
@not_implemented_for("multigraph")
|
|
def maximum_spanning_tree(G, weight="weight", algorithm="kruskal", ignore_nan=False):
|
|
"""Returns a maximum spanning tree or forest on an undirected graph `G`.
|
|
|
|
Parameters
|
|
----------
|
|
G : undirected graph
|
|
An undirected graph. If `G` is connected, then the algorithm finds a
|
|
spanning tree. Otherwise, a spanning forest is found.
|
|
|
|
weight : str
|
|
Data key to use for edge weights.
|
|
|
|
algorithm : string
|
|
The algorithm to use when finding a maximum spanning tree. Valid
|
|
choices are 'kruskal', 'prim', or 'boruvka'. The default is
|
|
'kruskal'.
|
|
|
|
ignore_nan : bool (default: False)
|
|
If a NaN is found as an edge weight normally an exception is raised.
|
|
If `ignore_nan is True` then that edge is ignored instead.
|
|
|
|
|
|
Returns
|
|
-------
|
|
G : EasyGraph Graph
|
|
A maximum spanning tree or forest.
|
|
|
|
|
|
Examples
|
|
--------
|
|
>>> G.add_edge(0, 3, weight=2)
|
|
>>> T = eg.maximum_spanning_tree(G)
|
|
>>> sorted(T.edges(data=True))
|
|
[(0, 1, {}), (0, 3, {'weight': 2}), (1, 2, {})]
|
|
|
|
|
|
Notes
|
|
-----
|
|
For Borůvka's algorithm, each edge must have a weight attribute, and
|
|
each edge weight must be distinct.
|
|
|
|
For the other algorithms, if the graph edges do not have a weight
|
|
attribute a default weight of 1 will be used.
|
|
|
|
There may be more than one tree with the same minimum or maximum weight.
|
|
See :mod:`easygraph.tree.recognition` for more detailed definitions.
|
|
|
|
Isolated nodes with self-loops are in the tree as edgeless isolated nodes.
|
|
|
|
"""
|
|
edges = list(
|
|
maximum_spanning_edges(G, algorithm, weight, data=True, ignore_nan=ignore_nan)
|
|
)
|
|
T = G.__class__() # Same graph class as G
|
|
for i in G.nodes:
|
|
T.add_node(i)
|
|
for i in edges:
|
|
(u, v, t) = i
|
|
T.add_edge(u, v, **t)
|
|
return T
|
|
|
|
|
|
def edge_boundary(G, nbunch1, nbunch2=None, data=False, default=None):
|
|
"""Returns the edge boundary of `nbunch1`.
|
|
|
|
The *edge boundary* of a set *S* with respect to a set *T* is the
|
|
set of edges (*u*, *v*) such that *u* is in *S* and *v* is in *T*.
|
|
If *T* is not specified, it is assumed to be the set of all nodes
|
|
not in *S*.
|
|
|
|
Parameters
|
|
----------
|
|
G : EasyGraph graph
|
|
|
|
nbunch1 : iterable
|
|
Iterable of nodes in the graph representing the set of nodes
|
|
whose edge boundary will be returned. (This is the set *S* from
|
|
the definition above.)
|
|
|
|
nbunch2 : iterable
|
|
Iterable of nodes representing the target (or "exterior") set of
|
|
nodes. (This is the set *T* from the definition above.) If not
|
|
specified, this is assumed to be the set of all nodes in `G`
|
|
not in `nbunch1`.
|
|
|
|
data : bool or object
|
|
This parameter has the same meaning as in
|
|
:meth:`MultiGraph.edges`.
|
|
|
|
default : object
|
|
This parameter has the same meaning as in
|
|
:meth:`MultiGraph.edges`.
|
|
|
|
Returns
|
|
-------
|
|
iterator
|
|
An iterator over the edges in the boundary of `nbunch1` with
|
|
respect to `nbunch2`. If `keys`, `data`, or `default`
|
|
are specified and `G` is a multigraph, then edges are returned
|
|
with keys and/or data, as in :meth:`MultiGraph.edges`.
|
|
|
|
Notes
|
|
-----
|
|
Any element of `nbunch` that is not in the graph `G` will be
|
|
ignored.
|
|
|
|
`nbunch1` and `nbunch2` are usually meant to be disjoint, but in
|
|
the interest of speed and generality, that is not required here.
|
|
|
|
"""
|
|
nset1 = {v for v in G if v in nbunch1}
|
|
# Here we create an iterator over edges incident to nodes in the set
|
|
# `nset1`. The `Graph.edges()` method does not provide a guarantee
|
|
# on the orientation of the edges, so our algorithm below must
|
|
# handle the case in which exactly one orientation, either (u, v) or
|
|
# (v, u), appears in this iterable.
|
|
edges = G.edges(nset1, data=data, default=default)
|
|
# If `nbunch2` is not provided, then it is assumed to be the set
|
|
# complement of `nbunch1`. For the sake of efficiency, this is
|
|
# implemented by using the `not in` operator, instead of by creating
|
|
# an additional set and using the `in` operator.
|
|
if nbunch2 is None:
|
|
return (e for e in edges if (e[0] in nset1) ^ (e[1] in nset1))
|
|
nset2 = set(nbunch2)
|
|
return (
|
|
e
|
|
for e in edges
|
|
if (e[0] in nset1 and e[1] in nset2) or (e[1] in nset1 and e[0] in nset2)
|
|
)
|
|
|
|
|
|
"""
|
|
Union-find data structure.
|
|
"""
|
|
|
|
|
|
class UnionFind:
|
|
"""Union-find data structure.
|
|
|
|
Each unionFind instance X maintains a family of disjoint sets of
|
|
hashable objects, supporting the following two methods:
|
|
|
|
- X[item] returns a name for the set containing the given item.
|
|
Each set is named by an arbitrarily-chosen one of its members; as
|
|
long as the set remains unchanged it will keep the same name. If
|
|
the item is not yet part of a set in X, a new singleton set is
|
|
created for it.
|
|
|
|
- X.union(item1, item2, ...) merges the sets containing each item
|
|
into a single larger set. If any item is not yet part of a set
|
|
in X, it is added to X as one of the members of the merged set.
|
|
|
|
Union-find data structure. Based on Josiah Carlson's code,
|
|
http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/215912
|
|
with significant additional changes by D. Eppstein.
|
|
http://www.ics.uci.edu/~eppstein/PADS/UnionFind.py
|
|
|
|
"""
|
|
|
|
def __init__(self, elements=None):
|
|
"""Create a new empty union-find structure.
|
|
|
|
If *elements* is an iterable, this structure will be initialized
|
|
with the discrete partition on the given set of elements.
|
|
|
|
"""
|
|
if elements is None:
|
|
elements = ()
|
|
self.parents = {}
|
|
self.weights = {}
|
|
for x in elements:
|
|
self.weights[x] = 1
|
|
self.parents[x] = x
|
|
|
|
def __getitem__(self, object):
|
|
"""Find and return the name of the set containing the object."""
|
|
|
|
# check for previously unknown object
|
|
if object not in self.parents:
|
|
self.parents[object] = object
|
|
self.weights[object] = 1
|
|
return object
|
|
|
|
# find basic of objects leading to the root
|
|
path = [object]
|
|
root = self.parents[object]
|
|
while root != path[-1]:
|
|
path.append(root)
|
|
root = self.parents[root]
|
|
|
|
# compress the basic and return
|
|
for ancestor in path:
|
|
self.parents[ancestor] = root
|
|
return root
|
|
|
|
def __iter__(self):
|
|
"""Iterate through all items ever found or unioned by this structure."""
|
|
return iter(self.parents)
|
|
|
|
def to_sets(self):
|
|
"""Iterates over the sets stored in this structure.
|
|
|
|
For example::
|
|
|
|
>>> partition = UnionFind("xyz")
|
|
>>> sorted(map(sorted, partition.to_sets()))
|
|
[['x'], ['y'], ['z']]
|
|
>>> partition.union("x", "y")
|
|
>>> sorted(map(sorted, partition.to_sets()))
|
|
[['x', 'y'], ['z']]
|
|
|
|
"""
|
|
# Ensure fully pruned paths
|
|
|
|
def groups(parents: dict):
|
|
sets = {}
|
|
for v, k in parents.items():
|
|
if k not in sets:
|
|
sets[k] = set()
|
|
sets[k].add(v)
|
|
return sets
|
|
|
|
for x in self.parents.keys():
|
|
_ = self[x] # Evaluated for side-effect only
|
|
|
|
yield from groups(self.parents).values()
|
|
|
|
def union(self, *objects):
|
|
"""Find the sets containing the objects and merge them all."""
|
|
# Find the heaviest root according to its weight.
|
|
roots = iter(sorted({self[x] for x in objects}, key=lambda r: self.weights[r]))
|
|
try:
|
|
root = next(roots)
|
|
except StopIteration:
|
|
return
|
|
|
|
for r in roots:
|
|
self.weights[root] += self.weights[r]
|
|
self.parents[r] = root
|