/
bst_construction.py
78 lines (73 loc) · 2.23 KB
/
bst_construction.py
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
class BST:
def __init__(self, val):
self.val = val
self.lc = None
self.rc = None
def insert(self, val):
cn = self
while True:
if val < cn.val:
if cn.lc is None:
cn.lc = BST(val)
break
else:
cn = cn.lc
else:
if cn.rc is None:
cn.rc = BST(val)
break
else:
cn = cn.rc
return self
def contains(self, val):
cn = self
while cn is not None:
if val < cn.val:
cn = cn.lc
elif val > cn.val:
cn = cn.rc
else:
return True
return False
def remove(self, val, pn = None):
cn = self
while cn is not None:
if val < cn.val:
pn = cn
cn = cn.lc
elif val > cn.val:
pn = cn
cn = cn.rc
else:
if cn.lc is not None and cn.rc is not None:
cn.val = cn.rc.get_min_val()
cn.rc.remove(cn.val, cn)
elif pn is None:
if cn.lc is not None:
cn.val = cn.lc.val
cn.rc = cn.lc.rc
cn.lc = cn.lc.lc
elif cn.rc is not None:
cn.val = cn.rc.val
cn.lc = cn.rc.lc
cn.rc = cn.rc.rc
else:
cn.val = None
elif pn.lc == cn:
pn.lc = cn.lc if cn.lc is not None else cn.rc
elif pn.rc == cn:
pn.rc = cn.lc if cn.lc is not None else cn.rc
break
return self
def get_min_val(self):
cn = self
while cn.lc is not None:
cn = cn.lc
return cn.val
bst = BST(50)
bst.insert(30)
bst.insert(70)
print(bst.contains(30)) # Should print True
print(bst.contains(40)) # Should print False
bst.remove(30)
print(bst.contains(30)) # Should print False