Skip to content

Commit

Permalink
Implement Kruskal traversal iterator
Browse files Browse the repository at this point in the history
Co-authored-by: kylebeggs <beggskw@gmail.com>
  • Loading branch information
Tortar and kylebeggs authored Jan 30, 2024
1 parent f6f8db6 commit 6c6b4a4
Showing 1 changed file with 46 additions and 0 deletions.
46 changes: 46 additions & 0 deletions src/iterators/kruskal.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# The kruskal algorithms are largely copied from Graphs/src/spanningtrees/kruskal.jl

struct KruskalIterator <: EdgeIterator
graph::AbstractGraph
connected_vs::IntDisjointSets
distmx::AbstractMatrix
edge_list

function KruskalIterator(graph, distmx=weights(g); minimize=true)
is_directed(graph) && throw(ArgumentError("Cannot use Kruskal on a directed graph."))
weights = Vector{eltype(distmx)}()
sizehint!(weights, ne(graph))
edge_list = collect(edges(graph))
for e in edge_list
push!(weights, distmx[src(e), dst(e)])
end
e = edge_list[sortperm(weights; rev=!minimize)]
new(graph, IntDisjointSets(nv(graph)), distmx, e)
end
end

Base.length(t::KruskalIterator) = nv(t.graph)-1
Base.eltype(t::KruskalIterator) = edgetype(t.graph)


"""
mutable struct KruskalIteratorState
`KruskalIteratorState` is a struct to hold the current state of iteration which is need for the Base.iterate() function.
"""
mutable struct KruskalIteratorState <: AbstractIteratorState
edge_id::Int
mst_len::Int
end


function Base.iterate(t::KruskalIterator, state::KruskalIteratorState=KruskalIteratorState(1,1))
while state.mst_len <= (nv(t.graph)-1)
i = state.edge_id
if !in_same_set(t.connected_vs, src(t.edge_list[i]), dst(t.edge_list[i]))
union!(t.connected_vs, src(t.edge_list[i]), dst(t.edge_list[i]))
state.mst_len += 1
return (t.edge_list[i], state)
end
state.edge_id += 1
end
end

0 comments on commit 6c6b4a4

Please sign in to comment.