diff --git a/astrodendro/plot.py b/astrodendro/plot.py
index 83de7dd..77bdd10 100644
--- a/astrodendro/plot.py
+++ b/astrodendro/plot.py
@@ -1,6 +1,7 @@
 # Licensed under an MIT open source license - see LICENSE
 
 import numpy as np
+from .structure import Structure
 
 
 class DendrogramPlotter(object):
@@ -190,7 +191,11 @@ def get_lines(self, structures=None, subtree=True, **kwargs):
         # Case 2: one structure is selected, and subtree is True
         else:
             if subtree:
-                if type(structures[0]) is int:
+                if isinstance(structures, int):
+                    structure = self.dendrogram[structures]
+                elif isinstance(structures, Structure):
+                    structure = structures
+                elif isinstance(structures[0], int):
                     structure = self.dendrogram[structures[0]]
                 else: 
                     structure = structures[0]