DeaD_EyE

iter_traverse

Jul 31st, 2020
1,471
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from collections import deque
  2. from operator import attrgetter
  3.  
  4.  
  5. class Node:
  6.     def __init__(self, name):
  7.         self.children = []
  8.         self.name = name
  9.  
  10.     def addChild(self, child):
  11.         self.children.append(child)
  12.  
  13.     def __repr__(self):
  14.         return "Node('{}')".format(self.name)
  15.  
  16.  
  17. def iter_traverse(node, attribute="children"):
  18.     seen = set()
  19.     stack = deque([node])
  20.     result = []
  21.     getter = attrgetter(attribute)
  22.     while stack:
  23.         node = stack.popleft()
  24.         node_id = id(node)
  25.         if node_id not in seen:
  26.             seen.add(node_id)
  27.             result.append(node)
  28.         children = getter(result[-1])
  29.         for child in children:
  30.             stack.append(child)
  31.     return result
  32.  
  33.  
  34. node1 = Node("one")
  35. node2 = Node("two")
  36. node3 = Node("three")
  37. [node3.addChild(Node("N{}".format(x))) for x in range(5)]
  38. node4 = Node("four")
  39. node1.addChild(node2)
  40. node2.addChild(node3)
  41. node3.addChild(node4)
  42. node3.addChild(node1) # CYCLE!
  43.  
  44. print(iter_traverse(node1))
  45. # [Node('one'), Node('two'), Node('three'), Node('N0'), Node('N1'), Node('N2'), Node('N3'), Node('N4'), Node('four')]
  46.  
RAW Paste Data