diff --git a/navis/morpho/subset.py b/navis/morpho/subset.py
index 82c5b438..39ed58a5 100644
--- a/navis/morpho/subset.py
+++ b/navis/morpho/subset.py
@@ -309,12 +309,8 @@ def _subset_treeneuron(x, subset, keep_disc_cn, prevent_fragments):
     )
 
     # Make sure any new roots or leafs are properly typed
-    # We won't produce new slabs but roots and leaves might change
-    x.nodes.loc[x.nodes.parent_id < 0, "type"] = "root"
-    x.nodes.loc[
-        (~x.nodes.node_id.isin(x.nodes.parent_id.values) & (x.nodes.parent_id >= 0)),
-        "type",
-    ] = "end"
+    # We won't produce new slabs but roots, branches and leaves might change
+    graph.classify_nodes(x, inplace=True)
 
     # Filter connectors
     if not keep_disc_cn and x.has_connectors: