@@ -164,15 +164,16 @@ def descendants(self, start, generations=None):
164164
165165 """
166166 visited = self .vertex_set ()
167+ visited .add (start )
167168 to_visit = deque ([(start , 0 )])
168169 while to_visit :
169170 vertex , depth = to_visit .popleft ()
170- visited . add ( vertex )
171- if generations is None or depth < generations :
172- to_visit . extend (
173- ( child , depth + 1 ) for child in self . children ( vertex )
174- if child not in visited
175- )
171+ if depth == generations :
172+ continue
173+ for child in self . children ( vertex ):
174+ if child not in visited :
175+ visited . add ( child )
176+ to_visit . append (( child , depth + 1 ) )
176177 return self .full_subgraph (visited )
177178
178179 def ancestors (self , start , generations = None ):
@@ -185,15 +186,16 @@ def ancestors(self, start, generations=None):
185186
186187 """
187188 visited = self .vertex_set ()
189+ visited .add (start )
188190 to_visit = deque ([(start , 0 )])
189191 while to_visit :
190192 vertex , depth = to_visit .popleft ()
191- visited . add ( vertex )
192- if generations is None or depth < generations :
193- to_visit . extend (
194- ( parent , depth + 1 ) for parent in self . parents ( vertex )
195- if parent not in visited
196- )
193+ if depth == generations :
194+ continue
195+ for parent in self . parents ( vertex ):
196+ if parent not in visited :
197+ visited . add ( parent )
198+ to_visit . append (( parent , depth + 1 ) )
197199 return self .full_subgraph (visited )
198200
199201 def shortest_path (self , start , end ):
0 commit comments