"""Independent implementation of a Trie tree."""
__all__ = ('Trie', 'TrieNode')
def _iterate_stringlike_objects(string):
for i in range(len(string)):
yield string[i:i + 1]
[docs]class Trie(object):
"""The object that manages the trie nodes."""
def __init__(self):
"""Initialize an empty trie."""
self.root = TrieNode(None, None)
def add(self, path, node_data):
"""Add the node data to the path described."""
node = self.root
for prefix in _iterate_stringlike_objects(path):
child = node.find_prefix(prefix)
if child is None:
child = node.add_child(prefix, [])
node = child
node.data.append(node_data)
def find(self, path):
"""Find a node based on the path provided."""
node = self.root
for prefix in _iterate_stringlike_objects(path):
child = node.find_prefix(prefix)
if child is None:
return None
node = child
return node
def traverse(self):
"""Traverse this tree.
This performs a depth-first pre-order traversal of children in this
tree. It returns the results consistently by first sorting the
children based on their prefix and then traversing them in
alphabetical order.
"""
return self.root.traverse()
class TrieNode(object):
"""The majority of the implementation details of a Trie."""
def __init__(self, prefix, data, children=None):
"""Initialize a TrieNode with data and children."""
self.children = children or {}
self.data = data
self.prefix = prefix
def __repr__(self):
"""Generate an easy to read representation of the node."""
return 'TrieNode(prefix={0}, data={1})'.format(
self.prefix, self.data
)
def find_prefix(self, prefix):
"""Find the prefix in the children of this node.
:returns: A child matching the prefix or None.
:rtype: :class:`~TrieNode` or None
"""
return self.children.get(prefix, None)
def add_child(self, prefix, data, children=None):
"""Create and add a new child node.
:returns: The newly created node
:rtype: :class:`~TrieNode`
"""
new_node = TrieNode(prefix, data, children)
self.children[prefix] = new_node
return new_node
def traverse(self):
"""Traverse children of this node.
This performs a depth-first pre-order traversal of the remaining
children in this sub-tree. It returns the results consistently by
first sorting the children based on their prefix and then traversing
them in alphabetical order.
"""
if not self.children:
return
for prefix in sorted(self.children):
child = self.children[prefix]
yield child
for child in child.traverse():
yield child