v

1
module datatypes
2

3
/// Internal representation of the tree node
4
@[heap]
5
struct BSTreeNode[T] {
6
mut:
7
	// Mark a node as initialized
8
	is_init bool
9
	// Value of the node
10
	value T
11
	// The parent of the node
12
	parent &BSTreeNode[T] = unsafe { 0 }
13
	// The left side with value less than the
14
	// value of this node
15
	left &BSTreeNode[T] = unsafe { 0 }
16
	// The right side with value grater than the
17
	// value of thiss node
18
	right &BSTreeNode[T] = unsafe { 0 }
19
}
20

21
// Create new root bst node
22
fn new_root_node[T](value T) &BSTreeNode[T] {
23
	return &BSTreeNode[T]{
24
		is_init: true
25
		value:   value
26
		parent:  new_none_node[T](true)
27
		left:    new_none_node[T](false)
28
		right:   new_none_node[T](false)
29
	}
30
}
31

32
// new_node creates a new bst node with a parent reference.
33
fn new_node[T](parent &BSTreeNode[T], value T) &BSTreeNode[T] {
34
	return &BSTreeNode[T]{
35
		is_init: true
36
		value:   value
37
		parent:  parent
38
	}
39
}
40

41
// new_none_node creates a dummy node.
42
fn new_none_node[T](init bool) &BSTreeNode[T] {
43
	return &BSTreeNode[T]{
44
		is_init: init
45
	}
46
}
47

48
// bind to an actual instance of a node.
49
fn (mut node BSTreeNode[T]) bind(mut to_bind BSTreeNode[T], left bool) {
50
	node.left = to_bind.left
51
	node.right = to_bind.right
52
	node.value = to_bind.value
53
	node.is_init = to_bind.is_init
54
	to_bind = new_none_node[T](false)
55
}
56

57
// Pure Binary Seach Tree implementation
58
//
59
// Pure V implementation of the Binary Search Tree
60
// Time complexity of main operation O(log N)
61
// Space complexity O(N)
62
pub struct BSTree[T] {
63
mut:
64
	root &BSTreeNode[T] = unsafe { 0 }
65
}
66

67
// insert give the possibility to insert an element in the BST.
68
pub fn (mut bst BSTree[T]) insert(value T) bool {
69
	if bst.is_empty() {
70
		bst.root = new_root_node(value)
71
		return true
72
	}
73
	return bst.insert_helper(mut bst.root, value)
74
}
75

76
// insert_helper walks the tree and inserts the given node.
77
fn (mut bst BSTree[T]) insert_helper(mut node BSTreeNode[T], value T) bool {
78
	if node.value < value {
79
		if unsafe { node.right != 0 } && node.right.is_init {
80
			return bst.insert_helper(mut node.right, value)
81
		}
82
		node.right = new_node(node, value)
83
		return true
84
	} else if node.value > value {
85
		if unsafe { node.left != 0 } && node.left.is_init {
86
			return bst.insert_helper(mut node.left, value)
87
		}
88
		node.left = new_node(node, value)
89
		return true
90
	}
91
	return false
92
}
93

94
// contains checks if an element with a given `value` is inside the BST.
95
pub fn (bst &BSTree[T]) contains(value T) bool {
96
	return bst.contains_helper(bst.root, value)
97
}
98

99
// contains_helper is a helper function to walk the tree, and return
100
// the absence or presence of the `value`.
101
fn (bst &BSTree[T]) contains_helper(node &BSTreeNode[T], value T) bool {
102
	if unsafe { node == 0 } || !node.is_init {
103
		return false
104
	}
105
	if node.value < value {
106
		return bst.contains_helper(node.right, value)
107
	} else if node.value > value {
108
		return bst.contains_helper(node.left, value)
109
	}
110
	assert node.value == value
111
	return true
112
}
113

114
// remove removes an element with `value` from the BST.
115
pub fn (mut bst BSTree[T]) remove(value T) bool {
116
	if bst.is_empty() {
117
		return false
118
	}
119
	return bst.remove_helper(mut bst.root, value, false)
120
}
121

122
fn (mut bst BSTree[T]) remove_helper(mut node BSTreeNode[T], value T, left bool) bool {
123
	if !node.is_init {
124
		return false
125
	}
126
	if node.value == value {
127
		if unsafe { node.left != 0 } && node.left.is_init {
128
			// In order to remove the element we need to bring up as parent the max of the
129
			// left sub-tree.
130
			mut max_node := bst.get_max_from_right(node.left)
131
			node.bind(mut max_node, true)
132
		} else if unsafe { node.right != 0 } && node.right.is_init {
133
			// Bring up the element with the minimum value in the right sub-tree.
134
			mut min_node := bst.get_min_from_left(node.right)
135
			node.bind(mut min_node, false)
136
		} else {
137
			mut parent := node.parent
138
			if left {
139
				parent.left = new_none_node[T](false)
140
			} else {
141
				parent.right = new_none_node[T](false)
142
			}
143
			node = new_none_node[T](false)
144
		}
145
		return true
146
	}
147

148
	if node.value < value {
149
		return bst.remove_helper(mut node.right, value, false)
150
	}
151
	return bst.remove_helper(mut node.left, value, true)
152
}
153

154
// get_max_from_right returns the max element of the BST following the right branch.
155
fn (bst &BSTree[T]) get_max_from_right(node &BSTreeNode[T]) &BSTreeNode[T] {
156
	if unsafe { node == 0 } {
157
		return new_none_node[T](false)
158
	}
159
	right_node := node.right
160
	if unsafe { right_node == 0 } || !right_node.is_init {
161
		return node
162
	}
163
	return bst.get_max_from_right(right_node)
164
}
165

166
// get_min_from_left returns the min element of the BST by following the left branch.
167
fn (bst &BSTree[T]) get_min_from_left(node &BSTreeNode[T]) &BSTreeNode[T] {
168
	if unsafe { node == 0 } {
169
		return new_none_node[T](false)
170
	}
171
	left_node := node.left
172
	if unsafe { left_node == 0 } || !left_node.is_init {
173
		return node
174
	}
175
	return bst.get_min_from_left(left_node)
176
}
177

178
// is_empty checks if the BST is empty
179
pub fn (bst &BSTree[T]) is_empty() bool {
180
	return unsafe { bst.root == 0 }
181
}
182

183
// in_order_traversal traverses the BST in order, and returns the result as an array.
184
pub fn (bst &BSTree[T]) in_order_traversal() []T {
185
	mut result := []T{}
186
	bst.in_order_traversal_helper(bst.root, mut result)
187
	return result
188
}
189

190
// in_order_traversal_helper helps traverse the BST, and accumulates the result in the `result` array.
191
fn (bst &BSTree[T]) in_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
192
	if unsafe { node == 0 } || !node.is_init {
193
		return
194
	}
195
	bst.in_order_traversal_helper(node.left, mut result)
196
	result << node.value
197
	bst.in_order_traversal_helper(node.right, mut result)
198
}
199

200
// post_order_traversal traverses the BST in post order, and returns the result in an array.
201
pub fn (bst &BSTree[T]) post_order_traversal() []T {
202
	mut result := []T{}
203
	bst.post_order_traversal_helper(bst.root, mut result)
204
	return result
205
}
206

207
// post_order_traversal_helper is a helper function that traverses the BST in post order,
208
// accumulating the result in an array.
209
fn (bst &BSTree[T]) post_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
210
	if unsafe { node == 0 } || !node.is_init {
211
		return
212
	}
213

214
	bst.post_order_traversal_helper(node.left, mut result)
215
	bst.post_order_traversal_helper(node.right, mut result)
216
	result << node.value
217
}
218

219
// pre_order_traversal traverses the BST in pre order, and returns the result as an array.
220
pub fn (bst &BSTree[T]) pre_order_traversal() []T {
221
	mut result := []T{}
222
	bst.pre_order_traversal_helper(bst.root, mut result)
223
	return result
224
}
225

226
// pre_order_traversal_helper is a helper function to traverse the BST
227
// in pre order and accumulates the results in an array.
228
fn (bst &BSTree[T]) pre_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
229
	if unsafe { node == 0 } || !node.is_init {
230
		return
231
	}
232
	result << node.value
233
	bst.pre_order_traversal_helper(node.left, mut result)
234
	bst.pre_order_traversal_helper(node.right, mut result)
235
}
236

237
// get_node is a helper method to ge the internal representation of the node with the `value`.
238
fn (bst &BSTree[T]) get_node(node &BSTreeNode[T], value T) &BSTreeNode[T] {
239
	if unsafe { node == 0 } || !node.is_init {
240
		return new_none_node[T](false)
241
	}
242
	if node.value == value {
243
		return node
244
	}
245

246
	if node.value < value {
247
		return bst.get_node(node.right, value)
248
	}
249
	return bst.get_node(node.left, value)
250
}
251

252
// to_left returns the value of the node to the left of the node with `value` specified if it exists,
253
// otherwise the a false value is returned.
254
//
255
// An example of usage can be the following one
256
//```v
257
// left_value, exist := bst.to_left(10)
258
//```
259
pub fn (bst &BSTree[T]) to_left(value T) !T {
260
	if bst.is_empty() {
261
		return error('BSTree is empty')
262
	}
263
	node := bst.get_node(bst.root, value)
264
	if !node.is_init {
265
		return error('BSTree is not initialized')
266
	}
267
	left_node := node.left
268
	return left_node.value
269
}
270

271
// to_right return the value of the element to the right of the node with `value` specified, if exist
272
// otherwise, the boolean value is false
273
// An example of usage can be the following one
274
//
275
//```v
276
// left_value, exist := bst.to_right(10)
277
//```
278
pub fn (bst &BSTree[T]) to_right(value T) !T {
279
	if bst.is_empty() {
280
		return error('BSTree is empty')
281
	}
282
	node := bst.get_node(bst.root, value)
283
	if !node.is_init {
284
		return error('BSTree is not initialized')
285
	}
286
	right_node := node.right
287
	return right_node.value
288
}
289

290
// max return the max element inside the BST.
291
// Time complexity O(N) if the BST is not balanced
292
pub fn (bst &BSTree[T]) max() !T {
293
	if bst.is_empty() {
294
		return error('BSTree is empty')
295
	}
296
	max := bst.get_max_from_right(bst.root)
297
	if !max.is_init {
298
		return error('BSTree is not initialized')
299
	}
300
	return max.value
301
}
302

303
// min return the minimum element in the BST.
304
// Time complexity O(N) if the BST is not balanced.
305
pub fn (bst &BSTree[T]) min() !T {
306
	if bst.is_empty() {
307
		return error('BSTree is empty')
308
	}
309
	min := bst.get_min_from_left(bst.root)
310
	if !min.is_init {
311
		return error('BSTree is not initialized')
312
	}
313
	return min.value
314
}
315