trie.py

from pprint import pprint

import types

def var_dump(obj):
  pprint(dump(obj))

def dump(obj):
  newobj = obj
  if isinstance(obj, list):
    newobj = []
    for item in obj:
      newobj.append(dump(item))
  elif isinstance(obj, tuple):
    temp = []
    for item in obj:
      temp.append(dump(item))
    newobj = tuple(temp)
  elif isinstance(obj, set):
    temp = []
    for item in obj:
      temp.append(str(dump(item)))
    newobj = set(temp)
  elif isinstance(obj, dict):
    newobj = {}
    for key, value in obj.items():
      newobj[str(dump(key))] = dump(value)
  elif isinstance(obj, types.FunctionType):
    newobj = repr(obj)
  elif '__dict__' in dir(obj):
    newobj = obj.__dict__.copy()
    if ' object at ' in str(obj) and not '__type__' in newobj:
      newobj['__type__']=str(obj).replace(" object at ", " #").replace("__main__.", "")
    for attr in newobj:
      newobj[attr]=dump(newobj[attr])
  return newobj

class Trie:
    def __init__(self, x = None):
        self.root = Trie.Node(None)
        self.leaf = x
        
    def search(self, seq):
        node = self.root
        for x in seq:
            node = node.get_child(x)
            if node is Node: return False
        # check leaf
        return node.get_child(self.leaf) is not None
    
    def insert(self, seq):
        node = self.root
        for x in seq:
            child = node.get_child(x)
            if not child:
                child = node.set_child(x)
            node = child
        # check leaf
        if not node.get_child(self.leaf):
            node.set_child(self.leaf)
            
    def delete(self, seq):
        node = self.root
        for x in req:
            node = node.get_child(x)
            if not node: return False
        # delete leaf
        return node.del_child(self.leaf)
        
    def traverse(self):
        node = self.root.child
        while node:
            for x in node.traverse(self.leaf):
                yield x
            node = node.bros
            
    def common_prefix(self, seq):
        node = self.root
        buff = []
        for x in seq:
            buff.append(x)
            node = node.get_child(x)
            if not node: return
        node = node.child
        while node:
            for x in node.traverse(self.leaf):
                yield buff + x
            node = node.bros

    class Node:
        def __init__(self, x, bros = None, child = None):
            self.data = x
            self.bros = bros
            self.child = child

        def set_child(self, x):
            child = Trie.Node(x, self.child)
            self.child = child
            return child

        def get_child(self, x):
            child = self.child
            while child:
                if child.data == x: break
                child = child.bros
            return child

        def del_child(self, x):
            child = self.child
            if child.data == x:
                self.child = child.bros
                return True
            else:
                while child.bros:
                    if child.bros.data == x:
                        child.bros = child.bros.bros
                        return True
                    child = child.bros
            return False

        def traverse(self, leaf):
            if self.data == leaf:
                yield []
            else:
                child = self.child
                while child:
                    for x in child.traverse(leaf):
                        yield [self.data] + x
                    child = child.bros
                    
if __name__ == '__main__':
    # suffix trie
    def make_suffix_trie(seq):
        a = Trie('LEAF')
        print('insert', seq[0:])
        a.insert(seq[0:])
        var_dump(a)
        print('insert', seq[0:-1])
        a.insert(seq[0:-1])
        var_dump(a)
        print('insert', seq[0:-1]+'x')
        a.insert(seq[0:-1]+'x')
        var_dump(a)
        return a
    
    s = make_suffix_trie('xyz')