Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added code to handle multi-graph in mst #7454

Merged
merged 7 commits into from
Jul 9, 2024
Merged

Added code to handle multi-graph in mst #7454

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6c4aee8
Added code to handle multi-graph in mst
Aditya-Shandilya1182 May 20, 2024
4dca079
Added code to handle multi-graph in mst
Aditya-Shandilya1182 May 20, 2024
f91dbc7
Apply Erik's suggested change
mjschwenne Jul 8, 2024
9653605
remove extra comment
mjschwenne Jul 8, 2024
74eb78d
add multigraph tests for SpanningTreeIterator
mjschwenne Jul 8, 2024
1181136
Merge branch 'main' into Aditya-Shandilya1182/main
mjschwenne Jul 8, 2024
f293e39
revert to nx.is_tree
mjschwenne Jul 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
31 changes: 23 additions & 8 deletions networkx/algorithms/tree/mst.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1125,19 +1125,34 @@ def _write_partition(self, partition):
A Partition dataclass describing a partition on the edges of the
graph.
"""
for u, v, d in self.G.edges(data=True):
if (u, v) in partition.partition_dict:
d[self.partition_key] = partition.partition_dict[(u, v)]
else:
d[self.partition_key] = EdgePartition.OPEN

partition_dict = partition.partition_dict
partition_key = self.partition_key
G = self.G

# Iterate over edges with for multi-graphs
mjschwenne marked this conversation as resolved.
Show resolved Hide resolved
if G.is_multigraph():
for u, v, k, d in G.edges(keys=True, data=True):
edge = (u, v, k)
d[partition_key] = partition_dict.get(edge, EdgePartition.OPEN)

# Iterate over edges for simple graphs
else:
for u, v, d in G.edges(data=True):
edge = (u, v)
d[partition_key] = partition_dict.get(edge, EdgePartition.OPEN)
mjschwenne marked this conversation as resolved.
Show resolved Hide resolved

def _clear_partition(self, G):
"""
Removes partition data from the graph
"""
for u, v, d in G.edges(data=True):
if self.partition_key in d:
del d[self.partition_key]
partition_key = self.partition_key
edges = (
G.edges(keys=True, data=True) if G.is_multigraph() else G.edges(data=True)
)
Comment on lines +1145 to +1147
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While we're here, I think this looks cleaner. Feel free to disregard.

Suggested change
edges = (
G.edges(keys=True, data=True) if G.is_multigraph() else G.edges(data=True)
)
edges = G.edges(keys=True, data=True) if G.is_multigraph() else G.edges(data=True)

for *e, d in edges:
Comment on lines +1145 to +1148
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this pattern! What do you think about using it above when setting d[partition_key]? This way, we wouldn't need separate for loops for if G.is_multigraph() and else.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looked neat this way.
Will it affect the performance of the code? I would love to hear your suggestions. I will make sure to change the code accordingly.

if partition_key in d:
del d[partition_key]


@nx._dispatchable(edge_attrs="weight")
Expand Down