Remove degree-zero vertices from steiner_tree output (#82)

ITensor · May 28, 2024 · a8535ec · a8535ec
1 parent 949d0a2
commit a8535ec
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/steiner_tree.jl b/src/steiner_tree.jl
@@ -9,5 +9,9 @@ using SimpleTraits: SimpleTraits, Not, @traitfn
     map(v -> vertex_positions(g)[v], term_vert),
     dist_matrix_to_position_dist_matrix(g, distmx),
   )
-  return typeof(g)(position_tree, map(v -> ordered_vertices(g)[v], vertices(position_tree)))
+  tree = typeof(g)(position_tree, map(v -> ordered_vertices(g)[v], vertices(position_tree)))
+  for v in copy(vertices(tree))
+    iszero(degree(tree, v)) && rem_vertex!(tree, v)
+  end
+  return tree
 end
diff --git a/test/test_namedgraph.jl b/test/test_namedgraph.jl
@@ -675,7 +675,8 @@ end
     st = steiner_tree(g, terminal_vertices)
     es = [(1, 2) => (1, 3), (1, 3) => (1, 4), (1, 4) => (2, 4), (2, 4) => (3, 4)]
     @test ne(st) == 4
-    @test nv(st) == 12
+    @test nv(st) == 5
+    @test !any(v -> iszero(degree(st, v)), vertices(st))
     for e in es
       @test has_edge(st, e)
     end