summaryrefslogtreecommitdiff
path: root/patricia.py
blob: cf34a8eaf48b39caf63766b1888703c1a5d40e61 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
"""A Python implementation of PATRICIA tree.

PATRICIA - Practical Algorithm to Retrieve Information Coded in Alphanumeric
           D.R.Morrison (1968).
See http://www.csse.monash.edu.au/~lloyd/tildeAlgDS/Tree/PATRICIA.html if you
want to know what's a PATRICIA tree...

TODO: _ advanced search
      _ profile code
      _ use mxTextTools ?

:copyright: 2000-2008 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
:contact: http://www.logilab.fr/ -- mailto:contact@logilab.fr
:license: General Public License version 2 - http://www.gnu.org/licenses
"""
__docformat__ = "restructuredtext en"

from warnings import warn
warn('logilab.common.patricia module is deprecated and will disappear in a near release',
     DeprecationWarning, stacklevel=2)

def prefix(prfx, string):
    """return the index of the first character from string which differs from
    prefix
    """
    i = 0
    while i < len(prfx):
        if i == len(string) or prfx[i] != string[i]:
            break
        i += 1
    return i

def split(index, string):
    """split a string on index, returning a 3-uple :
        (string before index, character at index, string after index)
    """
    return string[:index], string[index], string[index+1:]


class PatriciaNode:
    """a PATRICIA trie node
    """
    
    def __init__(self, value='', leaf=0, data=None):
        self.value = value
        self.edges = {}
        if leaf:
            self.datas = [data]
        else:
            self.datas = []
        
    def insert(self, string, data):
        """ insert the string in the trie and associate data to it
        if the string exists is the trie, data is added to the existing datas
        """
        # are we arrived ?
        if self.value == string:
            self.datas.append(data)
        # not yet !
        else:
            # check we don't break compression (value don't match)
            ind = prefix(self.value, string)
            if ind < len(self.value):
                # split this node
                pfx, e, self.value = split(ind, self.value)
                if ind < len(string):
                    n = PatriciaNode(pfx)
                    n.edges[string[ind]] = PatriciaNode(string[ind+1:], 1, data)
                else:
                    n = PatriciaNode(pfx, 1, data)
                n.edges[e] = self
                return n
            n_pfx, n_e, n_sfx = split(len(self.value), string)
            if n_e in self.edges:
                self.edges[n_e] = self.edges[n_e].insert(n_sfx, data)
            else:
                self.edges[n_e] = PatriciaNode(n_sfx, 1, data)
        return self

    def remove(self, string):
        """ return datas associated with string and remove string from the trie
        raise KeyError if the key isn't found
        FIXME: we should change the trie structure
        """
        if string == self.value and self.datas:
            datas = self.datas
            self.datas = []
            return datas
        else: 
            pfx, e, sfx = split(len(self.value), string)
            if self.value == pfx:
                return self.edges[e].remove(sfx)
        raise KeyError(string)
    
    def lookup(self, string):
        """ return datas associated with string
        raise KeyError if the key isn't found
        """
        if string == self.value:
            if self.datas:
                return self.datas
            raise KeyError(string)
        else: # len(self.value) < len(string): 
            pfx, e, sfx = split(len(self.value), string)
            if self.value == pfx:
                return self.edges[e].lookup(sfx)
        raise KeyError(string)
    
    def pfx_search(self, pfx, depth=-1):
        """ return all string with prefix pfx """
        sfxs = []
        if pfx and self.value[:len(pfx)] != pfx:
            pfx, e, sfx = split(len(self.value), pfx)
            if self.value == pfx and  e in self.edges:
                sfxs = ['%s%s%s' % (self.value, e, sfx)
                        for sfx in self.edges[e].pfx_search(sfx, depth)]
        else:
            if depth != 0:
                for e, child in self.edges.items():
                    search = child.pfx_search('', depth-1-len(self.value))
                    sfxs += ['%s%s%s' % (self.value, e, sfx)
                             for sfx in search]
            if (depth < 0 or len(self.value) <= depth):
                if self.datas:
                    sfxs.append(self.value)
        return sfxs
        
    def __str__(self, indent=''):
        node_str = ''.join([' %s%s:\n%s' % (indent, key,
                                            a.__str__('  %s' % indent))
                            for key, a in self.edges.items()])
        return '%s%s, %s\n%s' % (indent, self.value, self.datas, node_str)

    def __repr__(self):
        return '<PatriciaNode id=%s value=%s childs=%s datas=%s>' % (
            id(self), self.value, self.edges.keys(), self.datas)


class PatriciaTrie:
    """ wrapper class for a patricia tree
    delegates to the root of the tree (PatriciaNode)
    """
    
    def __init__(self):
        self._trie = None
        self.words = 0

    def insert(self, string, data=None):
        """ insert a string into the tree """
        self.words += 1
        if self._trie is None:
            self._trie = PatriciaNode(string, 1, data)
        else:
            self._trie = self._trie.insert(string, data)
            
    def remove(self, string):
        """ remove a string from the tree """
        if self._trie is not None:
            return self._trie.remove(string)
        raise KeyError(string)

    def lookup(self, string):
        """ look for a string into the tree """
        if self._trie is not None:
            return self._trie.lookup(string)
        raise KeyError(string)

    def pfx_search(self, string, depth=-1):
        """ search all words begining by <string> """
        if self._trie is not None:
            return self._trie.pfx_search(string, depth)
        raise KeyError(string)

    def __str__(self):
        return self._trie.__str__()
    
    def __repr__(self):
        return '<PatriciaTrie id=%s words=%s>' % (id(self), self.words)