class BinaryTreeSearch:
def __init__(self, fileName):
self.tree = BinaryTree()
with open(fileName) as f:
flag = False
lastWord = ''
preLastWord = ''
oneWordLineWord = ''
for line in f:
words = re.split(r'\W+|\d+|_+', line)
words = list(filter(bool, words))
if len(words) == 0: continue
if oneWordLineWord != '':
words.insert(0, oneWordLineWord)
oneWordLineWord = ''
if len(words) == 1:
oneWordLineWord = words[0].lower()
continue
if flag:
#pre last first
node = self.tree.find(
NodeWithData.makeKey(preLastWord, lastWord))
if node == None:
node = NodeWithData(preLastWord, lastWord)
self.tree.add(node)
node.UpdateDict(preLastWord, lastWord, words[0].lower())
#last first second
node = self.tree.find(
NodeWithData.makeKey(lastWord, words[0].lower()))
if node == None:
node = NodeWithData(lastWord, words[0].lower())
self.tree.add(node)
node.UpdateDict(lastWord, words[0].lower(),
words[1].lower())
else:
flag = True
lastWord = words[-1].lower()
preLastWord = words[-2].lower()
for i in range(len(words) - 2):
w1 = words[i].lower()
w2 = words[i + 1].lower()
node = self.tree.find(NodeWithData.makeKey(w1, w2))
if node == None:
node = NodeWithData(w1, w2)
self.tree.add(node)
node.UpdateDict(w1, w2, words[i + 2].lower())
if oneWordLineWord + preLastWord + lastWord != '':
node = self.tree.find(
NodeWithData.makeKey(preLastWord, lastWord))
if node == None:
node = NodeWithData(preLastWord, lastWord)
self.tree.add(node)
node.UpdateDict(preLastWord, lastWord, oneWordLineWord)
def search(self, first, second):
first, second = first.lower(), second.lower()
key = NodeWithData.makeKey(first, second)
node = self.tree.find(key)
if node == None: return []
return node.GetTrigrams(first, second)
def test_delete(self):
var = globalVar() #Contains variables for tree insertion
bTree = BinaryTree()
bTree.insert(value= var.p1)
bTree.insert(value= var.p2)
bTree.insert(value= var.p3)
self.assertEqual(bTree.find(var.p3).value, var.p3)
bTree.delete(var.p3)
with self.assertRaises(Exception):
bTree.find(var.p3).value#Looking for deleted element
def test_find(self):
var = globalVar() #Contains variables for tree insertion
bTree = BinaryTree()
bTree.insert(value= var.p1)
bTree.insert(value= var.p2)
bTree.insert(value= var.p3)
bTree.insert(value= var.p5)
bTree.insert(value= var.p4)
self.assertEqual(bTree.find(var.p1).value, var.p1)
self.assertEqual(bTree.find(var.p4).value, var.p4)
def test_features(self):
bt = BinaryTree(10)
bt.add(6)
bt.add(18)
bt.add(4)
bt.add(8)
bt.add(15)
bt.add(21)
self.assertEqual(bt.count(), 7)
self.assertEqual(bt.min(), 4)
self.assertEqual(bt.max(), 21)
self.assertEqual(bt.has(8), True)
self.assertEqual(bt.has(11), False)
self.assertNotEqual(bt.find(15), None)
self.assertEqual(bt.find(11), None)
bt.delete(10)
self.assertEqual(bt.root.data, 15)
self.assertEqual(bt.find(10), None)
self.assertEqual(bt.count(), 6)
self.assertEqual([x for x in bt.items()], [4, 6, 8, 15, 18, 21])
self.assertEqual([x for x in bt.items('pre')], [15, 6, 4, 8, 18, 21])
self.assertEqual([x for x in bt.items('pos')], [4, 8, 6, 21, 18, 15])
self.assertEqual([x for x in bt.items('lev')], [15, 6, 18, 4, 8, 21])
