Commit 70c13294 authored by Jeroen F.J. Laros's avatar Jeroen F.J. Laros

Added counting.

parent c1c3d4ed
import itertools
def _add(root, word):
def _add(root, word, count):
"""Add a word to the trie.
:arg dict root: Root of the trie.
:arg str word: A word.
:arg int count: Multiplicity of `word`.
"""
node = root
......@@ -14,7 +15,9 @@ def _add(root, word):
node[char] = {}
node = node[char]
node[''] = {}
if '' not in node:
node[''] = {'count': 0}
node['']['count'] += count
def _find(root, word):
......@@ -35,16 +38,19 @@ def _find(root, word):
return node
def _remove(node, word):
def _remove(node, word, count):
"""Remove a word from a trie.
:arg dict node: Current node.
:arg str word: Word to be removed.
:arg int count: Multiplicity of `word`, force remove if this is -1.
:returns bool:
:returns bool: True if the last occurrence of `word` is removed.
"""
if not word:
if '' in node:
node['']['count'] -= count
if node['']['count'] < 1 or count == -1:
node.pop('')
return True
return False
......@@ -53,7 +59,7 @@ def _remove(node, word):
if car not in node:
return False
result = _remove(node[car], cdr)
result = _remove(node[car], cdr, count)
if result:
if not node[car]:
node.pop(car)
......@@ -194,11 +200,17 @@ class Trie(object):
def __iter__(self):
return _iterate('', self.root)
def add(self, word):
_add(self.root, word)
def add(self, word, count=1):
_add(self.root, word, count)
def get(self, word):
node = _find(self.root, word)
if '' in node:
return node['']
return None
def remove(self, word):
return _remove(self.root, word)
def remove(self, word, count=1):
return _remove(self.root, word, count)
def has_prefix(self, word):
return _find(self.root, word) != {}
......
......@@ -8,7 +8,7 @@ from dict_trie import Trie
class TestTrie(object):
def setup(self):
self._trie = Trie(['abc', 'abd', 'test', 'te'])
self._trie = Trie(['abc', 'abd', 'abd', 'test', 'te'])
def test_empty(self):
assert Trie().root == {}
......@@ -17,11 +17,11 @@ class TestTrie(object):
assert self._trie.root == {
'a': {
'b': {
'c': {'': {}},
'd': {'': {}}}},
'c': {'': {'count': 1}},
'd': {'': {'count': 2}}}},
't': {'e': {
'': {},
's': {'t': {'': {}}}}}}
'': {'count': 1},
's': {'t': {'': {'count': 1}}}}}}
def test_word_present(self):
assert 'abc' in self._trie
......@@ -60,6 +60,20 @@ class TestTrie(object):
self._trie.add('abx')
assert 'abx' in self._trie
def test_get_present(self):
assert self._trie.get('abc')['count'] == 1
def test_get_absent(self):
assert not self._trie.get('abx')
def test_add_twice(self):
self._trie.add('abc')
assert self._trie.get('abc')['count'] == 2
def test_add_multiple(self):
self._trie.add('abc', 2)
assert self._trie.get('abc')['count'] == 3
def test_remove_present(self):
assert self._trie.remove('test')
assert 'test' not in self._trie
......@@ -76,6 +90,23 @@ class TestTrie(object):
def test_remove_prefix_absent(self):
assert not self._trie.remove('ab')
def test_remove_twice(self):
self._trie.add('abc')
assert not self._trie.remove('abc')
assert self._trie.get('abc')['count'] == 1
assert self._trie.remove('abc')
assert 'abc' not in self._trie
def test_remove_multile(self):
self._trie.add('abc', 3)
assert not self._trie.remove('abc', 2)
assert self._trie.get('abc')['count'] == 2
def test_remove_force(self):
self._trie.add('abc')
assert self._trie.remove('abc', -1)
assert 'abc' not in self._trie
def test_iter(self):
assert list(self._trie) == ['abc', 'abd', 'te', 'test']
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment