import networkx as nx
import itertools
from typing import Iterable
# We use a module-level constant for inf to get around an issue
# where flake8 will complain about `float('inf')` as a default arg
_INF = float('inf')
def _add_edge_with_data(from_graph, to_graph, edge):
to_graph.add_edge(*edge)
to_graph[edge[0]][edge[1]].update(from_graph.get_edge_data(
edge[0], edge[1]))
[docs]def get_subgraphs(g: nx.Graph,
root,
*,
node_weight_map=None,
max_weight: float = _INF) -> Iterable[nx.Graph]:
"""
Return a generator over all subgraphs of `g` weighing at most `max_weight`
and with root `root`
For every iteration of this generator, a `nx.Graph` will be yielded that
is a subgraph of `g`. The same `nx.Graph` is used for every iteration so
if you need to save the state of the subgraph, make a copy of it
(`subgraph.copy()`).
All nodes are given a default weight of 1. If `node_weight_map` is supplied,
weights of nodes are determined by looking them up in the map. If they
are not found in the map, the default of 1 is used.
"""
for subtree in get_subtrees(g,
root,
node_weight_map=node_weight_map,
max_weight=max_weight):
yield subtree
cyclic_edges = []
for e in g.edges:
if e not in subtree.edges and e[0] in subtree.nodes and e[
1] in subtree.nodes:
cyclic_edges.append(e)
for n in range(1, len(cyclic_edges) + 1):
for edges in itertools.combinations(cyclic_edges, n):
subtree.add_edges_from(edges)
for e in edges:
_add_edge_with_data(g, subtree, e)
yield subtree
subtree.remove_edges_from(edges)
[docs]def get_subtrees(g: nx.Graph,
root,
*,
node_weight_map=None,
max_weight: float = _INF) -> Iterable[nx.Graph]:
"""
Return a generator over all subtrees of `g` weighing at most `max_weight`
and with root `root`
For every iteration of this generator, a `nx.Graph` will be yielded that
is a subtree of `g`. The same `nx.Graph` is used for every iteration so
if you need to save the state of the subtree, make a copy of it
(`subtree.copy()`).
All nodes are given a default weight of 1. If `node_weight_map` is supplied,
weights of nodes are determined by looking them up in the map. If they
are not found in the map, the default of 1 is used.
"""
if node_weight_map is None:
node_weight_map = dict()
current_subtree = nx.Graph()
current_subtree.add_nodes_from([(root, g.nodes[root])])
root_weight = node_weight_map.get(root, 1)
if max_weight < root_weight:
raise ValueError("max_weight must be larger than the weight of the "
"root node.")
remaining_weight = max_weight - root_weight
for _ in _subtrees(g,
current_subtree,
root,
node_weight_map,
max_weight=remaining_weight):
yield current_subtree
def _subtrees(g: nx.Graph,
subg: nx.Graph,
root,
node_weight_map,
max_weight=_INF):
"""
Generator that will modify the subgraph `subg` into a new subtree of
graph `g` on every iteration.
"""
assert root in subg
yield max_weight
unexplored_n = []
for n in g.neighbors(root):
if n not in subg and node_weight_map.get(n, 1) <= max_weight:
unexplored_n.append(n)
if not unexplored_n:
return
# For each combination of unexplored neighbors, add the edges to
# the neighbors and then iterate over all subtrees rooted in those
# neighbors.
for i in range(1, len(unexplored_n) + 1):
for n_to_explore in itertools.combinations(unexplored_n, i):
edges = [[n, root] for n in n_to_explore]
weight_of_n = sum(node_weight_map.get(n, 1) for n in n_to_explore)
if weight_of_n > max_weight:
continue
remaining_weight = max_weight - weight_of_n
for n in n_to_explore:
subg.add_nodes_from([(n, g.nodes[n])])
for e in edges:
_add_edge_with_data(g, subg, e)
yield from _product_of_subtrees(g, subg, n_to_explore,
node_weight_map, remaining_weight)
for e in edges:
subg.remove_edge(*e)
for n in n_to_explore:
subg.remove_node(n)
def _product_of_subtrees(g, subg, roots, node_weight_map, max_weight):
"""
Generator over the product of subtrees rooted with `roots` while
respecting `node_weight_map` and `max_weight`.
The iteration order advances subtrees from right to left (similar to
an odometer). The subtree rooted in the rightmost root is advanced
on every iteration.
"""
if len(roots) == 1:
yield from _subtrees(g, subg, roots[-1], node_weight_map, max_weight)
else:
assert roots
for remaining_wt in _product_of_subtrees(g, subg, roots[:-1],
node_weight_map, max_weight):
yield from _subtrees(g, subg, roots[-1], node_weight_map,
remaining_wt)