def delete(self, word):
# If tree is empty
if not self.root:
return False
# If trie is not empty
# Put the nodes in a stack
current = self.root
node_stack = Stack()
for char in word:
if char not in current.children:
return False
current = current.children[char]
node_stack.push(current)
if not current.is_end:
return False
# Delete node from the end of key if node does not have children
current.is_end = False
for i in range(node_stack.count()):
current = node_stack.pop()
prev = node_stack.peek()
if not current.children and prev:
del prev.children[current.char]
return True
def test_peek_empty(self):
mystack = Stack()
result = mystack.peek()
self.assertEqual(result, None)
self.assertEqual(mystack.count(), 0)
def test_push(self):
mystack = Stack()
mystack.push('Python0')
self.assertEqual(mystack.count(), 1)
self.assertEqual(mystack.peek(), 'Python0')