Post Order N-Ary Tree Traversal using Iteration

sunglasses.py

''' Question: N-Ary tree post order traversal
	Twist: Using only iteration!

	Rationale:
	This essentially becomes a Dynamic Programming question. Instead
	of using Top-Down recursion. We build up our solution from the bottom
	up using a list and our base case. The base case is if a node has no
	children.

	I built a simply NaryTree for testing and a function that takes a node
	to find the Post Order traversal for that sub tree. 
'''

class NaryTree(object):
	''' N-Ary Tree that has a root node and size '''

	class Node(object):
		''' Node object that contains
			a data field and list of children '''
		def __init__(self, data):
			self.data = data
			self.children = []

		def __str__(self):
			''' For debugging '''
			children = str([str(c.data) for c in self.children])
			return 'Node '+str(self.data)+': Children: '+children

		def add_child(self, data):
			''' Adds a child and returns that child '''
			child = NaryTree.Node(data)
			self.children.append(child)
			return child

		def has_children(self):
			''' Helper for dpPostOrder '''
			return len(self.children) != 0

	def __init__(self):
		self.root = None
		self.size = 0

	def __len__(self):
		return self.size

	def add(self, data, node=None):
		if self.root is None:
			self.root = NaryTree.Node(data)
			self.size += 1
			return self.root

		if node is None:
			node = self.root

		child = node.add_child(data)
		self.size += 1
		return child


def main():
	tree = NaryTree()

	A = tree.add('A')
	B = tree.add('B', A)
	C = tree.add('C', A)
	D = tree.add('D', A)
	E = tree.add('E', B)
	F = tree.add('F', B)

	print(A)
	print(B)
	print(str([str(node.data) for node in dpPostOrder(tree.root)]))

def dpPostOrder(node):
	result = []
	remaining = []
	remaining.append(node)
	memo = {}

	while len(remaining) != 0:
		cur = remaining[-1]

		if not cur.has_children() or cur in memo:
			result.append(remaining.pop())
		else: 
			for child in cur.children[::-1]:
				print('On node:'+str(cur.data)+' Adding node:'+str(child.data))
				remaining.append(child)
			memo[cur] = True
	return result

if __name__ == '__main__':
	main()