├── fiby_tree ├── __init__.mojo └── fiby_tree.mojo ├── trie_dict ├── __init__.mojo └── dict.mojo ├── helpers ├── __init__.mojo └── helpers.mojo ├── README.md ├── left_child_right_sibling ├── __init__.mojo ├── lcrs_tree_builder.mojo └── lcrs_tree.mojo ├── trie_dict_sample.mojo ├── LICENSE ├── fiby_tree_string_benachmarks.mojo ├── trie_dict_benchmarks.mojo ├── lcrs_tree_sample.mojo ├── fiby_tree_int_benchmarks.mojo └── fiby_tree_int_tests.mojo /fiby_tree/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .fiby_tree import FibyTree -------------------------------------------------------------------------------- /trie_dict/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .dict import TrieDict 2 | -------------------------------------------------------------------------------- /helpers/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .helpers import cmp_strl, stsl, int_to_str -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # mojo-trees 2 | Experimental Tree data structures in Mojo 3 | -------------------------------------------------------------------------------- /left_child_right_sibling/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .lcrs_tree import LCRSTree 2 | from .lcrs_tree_builder import LCRSTreeBuilder -------------------------------------------------------------------------------- /helpers/helpers.mojo: -------------------------------------------------------------------------------- 1 | from math import min 2 | 3 | fn int_to_str(i: Int) -> String: 4 | return String(i) 5 | 6 | fn int_eq(i1: Int, i2: Int) -> Bool: 7 | return i1 == i2 8 | 9 | fn cmp_strl(a: StringLiteral, b: StringLiteral) -> Int: 10 | let l = min(len(a), len(b)) 11 | let p1 = DTypePointer[DType.int8](a.data()).bitcast[DType.uint8]() 12 | let p2 = DTypePointer[DType.int8](b.data()).bitcast[DType.uint8]() 13 | let diff = memcmp(p1, p2, l) 14 | 15 | return diff if diff != 0 else len(a) - len(b) 16 | 17 | fn stsl(a: StringLiteral) -> String: 18 | return a -------------------------------------------------------------------------------- /trie_dict_sample.mojo: -------------------------------------------------------------------------------- 1 | from trie_dict import TrieDict 2 | 3 | 4 | fn main(): 5 | var t = TrieDict[Int]() 6 | _ = t.put("Maxim", 12) 7 | t.debug() 8 | _ = t.put("Max", 11) 9 | t.debug() 10 | _ = t.put("Marina", 13) 11 | t.debug() 12 | _ = t.put("Marinala", 14) 13 | t.debug() 14 | _ = t.put("Leo", 15) 15 | t.debug() 16 | _ = t.put("Daria", 16) 17 | t.debug() 18 | _ = t.put("Dario", 17) 19 | t.debug() 20 | _ = t.put("Dominique", 18) 21 | t.debug() 22 | 23 | print("Marina", t.__contains__("Marina")) 24 | print("Mari", t.__contains__("Mari")) 25 | print("Dominique", t.__contains__("Dominique")) 26 | 27 | print("Dominique", t.get("Dominique", 0)) 28 | print("Dom", t.get("Dom", 0)) 29 | print("Daria", t.get("Daria", 0)) 30 | 31 | _ = t.put("Daria", 26) 32 | print("Daria", t.get("Daria", 0)) 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Maxim Zaks 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /left_child_right_sibling/lcrs_tree_builder.mojo: -------------------------------------------------------------------------------- 1 | from .lcrs_tree import LCRSTree 2 | 3 | struct LCRSTreeBuilder[T: AnyType]: 4 | var _tree: LCRSTree[T] 5 | var _queue: DynamicVector[UInt16] 6 | 7 | fn __init__(inout self, root: T): 8 | self._tree = LCRSTree(root) 9 | self._queue = DynamicVector[UInt16]() 10 | self._queue.push_back(0) 11 | 12 | fn __moveinit__(inout self, owned existing: Self): 13 | self._tree = existing._tree^ 14 | self._queue = existing._queue 15 | 16 | fn node(owned self, node: T) -> LCRSTreeBuilder[T]: 17 | let last = self._queue[len(self._queue) - 1] 18 | self._queue.push_back(self._tree.add_child(node, last)) 19 | return self^ 20 | 21 | fn leaf(owned self, node: T) -> LCRSTreeBuilder[T]: 22 | let last = self._queue[len(self._queue) - 1] 23 | _ = self._tree.add_child(node, last) 24 | return self^ 25 | 26 | fn up(owned self) -> LCRSTreeBuilder[T]: 27 | _ = self._queue.pop_back() 28 | return self^ 29 | 30 | fn tree(owned self) -> LCRSTree[T]: 31 | return self._tree^ -------------------------------------------------------------------------------- /fiby_tree_string_benachmarks.mojo: -------------------------------------------------------------------------------- 1 | from fiby_tree import FibyTree 2 | from helpers import cmp_strl, stsl 3 | from time import now 4 | from math import min 5 | 6 | 7 | fn main(): 8 | var tik = now() 9 | var ft = FibyTree[StringLiteral, cmp_strl, stsl]('Lorem', 'ipsum', 'dolor', 'sit', 'amet,', 'consectetur', 'adipiscing', 'elit.', 'Quisque', 'orci', 'urna,', 'pretium', 'et', 'porta', 'ac,', 'porttitor', 'sit', 'amet', 'sem.', 'Fusce', 'sagittis', 'lorem', 'neque,', 'vitae', 'sollicitudin', 'elit', 'suscipit', 'et.', 'In', 'interdum', 'convallis', 'nisl', 'in', 'ornare.', 'Vestibulum', 'ante', 'ipsum', 'primis', 'in', 'faucibus', 'orci', 'luctus', 'et', 'ultrices', 'posuere', 'cubilia', 'curae;', 'Aliquam', 'erat', 'volutpat.', 'Morbi', 'mollis', 'iaculis', 'lectus', 'ac', 'tincidunt.', 'Fusce', 'nisi', 'lacus,', 'semper', 'eu', 'dignissim', 'et,', 'malesuada', 'non', 'mi.', 'Sed', 'euismod', 'urna', 'vel', 'elit', 'faucibus,', 'eu', 'bibendum', 'ante', 'fringilla.', 'Curabitur', 'tempus', 'in', 'turpis', 'at', 'mattis.', 'Aliquam', 'erat', 'volutpat.', 'Donec', 'maximus', 'elementum', 'felis,', 'sit', 'amet', 'dignissim', 'augue', 'tincidunt', 'blandit.', 'Aliquam', 'fermentum,', 'est', 'eu', 'mollis.') 10 | var tok = now() 11 | var duration = tok - tik 12 | print("Create 100 entry set in", duration, "ns") 13 | print("Set len:", ft.__len__()) 14 | 15 | tik = now() 16 | ft.balance() 17 | tok = now() 18 | duration = tok - tik 19 | print("Balanced the set in:", duration, "ns") 20 | 21 | tik = now() 22 | let elements = ft.sorted_elements() 23 | tok = now() 24 | duration = tok - tik 25 | print("Get sorted elements in:", duration, "ns") 26 | 27 | var total = 0 28 | for i in range(len(elements)): 29 | tik = now() 30 | let r = ft.__contains__(elements[i]) 31 | tok = now() 32 | total += tok - tik 33 | 34 | print("Check contains in:", total / len(elements), "ns on avg per element") 35 | 36 | total = 0 37 | for i in range(len(elements)): 38 | tik = now() 39 | let r = ft.delete(elements[i]) 40 | tok = now() 41 | total += tok - tik 42 | 43 | print("Delete all in:", total / len(elements), "ns on avg per element") 44 | -------------------------------------------------------------------------------- /trie_dict_benchmarks.mojo: -------------------------------------------------------------------------------- 1 | from time import now 2 | from trie_dict import TrieDict 3 | 4 | fn vec[T: AnyType](*elements: T) -> DynamicVector[T]: 5 | let elements_list: VariadicList[T] = elements 6 | var result = DynamicVector[T](len(elements_list)) 7 | for i in range(len(elements_list)): 8 | result.push_back(elements[i]) 9 | return result 10 | 11 | fn main(): 12 | let corpus = vec('Lorem', 'ipsum', 'dolor', 'sit', 'amet,', 'consectetur', 'adipiscing', 'elit.', 'Quisque', 'orci', 'urna,', 'pretium', 'et', 'porta', 'ac,', 'porttitor', 'sit', 'amet', 'sem.', 'Fusce', 'sagittis', 'lorem', 'neque,', 'vitae', 'sollicitudin', 'elit', 'suscipit', 'et.', 'In', 'interdum', 'convallis', 'nisl', 'in', 'ornare.', 'Vestibulum', 'ante', 'ipsum', 'primis', 'in', 'faucibus', 'orci', 'luctus', 'et', 'ultrices', 'posuere', 'cubilia', 'curae;', 'Aliquam', 'erat', 'volutpat.', 'Morbi', 'mollis', 'iaculis', 'lectus', 'ac', 'tincidunt.', 'Fusce', 'nisi', 'lacus,', 'semper', 'eu', 'dignissim', 'et,', 'malesuada', 'non', 'mi.', 'Sed', 'euismod', 'urna', 'vel', 'elit', 'faucibus,', 'eu', 'bibendum', 'ante', 'fringilla.', 'Curabitur', 'tempus', 'in', 'turpis', 'at', 'mattis.', 'Aliquam', 'erat', 'volutpat.', 'Donec', 'maximus', 'elementum', 'felis,', 'sit', 'amet', 'dignissim', 'augue', 'tincidunt', 'blandit.', 'Aliquam', 'fermentum,', 'est', 'eu', 'mollis.') 13 | 14 | var t = TrieDict[Int]() 15 | var tik = now() 16 | var tok = now() 17 | var total = 0 18 | for i in range(len(corpus)): 19 | tik = now() 20 | t.put(corpus[i], i) 21 | tok = now() 22 | total += tok - tik 23 | 24 | print("Add 100 elements in", total / len(corpus), "ns on avg per entry") 25 | 26 | for i in range(len(corpus)): 27 | tik = now() 28 | let r = t.__contains__(corpus[i]) 29 | tok = now() 30 | total += tok - tik 31 | 32 | print("Lookup 100 elements in", total / len(corpus), "ns on avg per entry") 33 | 34 | for i in range(len(corpus)): 35 | tik = now() 36 | let r = t.delete(corpus[i]) 37 | tok = now() 38 | total += tok - tik 39 | 40 | print("Delete 100 elements in", total / len(corpus), "ns on avg per entry") 41 | print("Dict len:", t.__len__()) 42 | -------------------------------------------------------------------------------- /lcrs_tree_sample.mojo: -------------------------------------------------------------------------------- 1 | from left_child_right_sibling import LCRSTree, LCRSTreeBuilder 2 | from helpers import int_to_str 3 | 4 | fn v1(i: UInt16, e: Int) -> Bool: 5 | print(i, e) 6 | return True 7 | 8 | fn main(): 9 | var t = LCRSTree(12) 10 | let n1 = t.add_child(15) 11 | let n2 = t.add_child(24) 12 | let n3 = t.add_child(75) 13 | _ = t.add_child(88) 14 | 15 | _ = t.add_child(45, n1) 16 | _ = t.add_child(55, n1) 17 | 18 | _ = t.add_child(590, n2) 19 | _ = t.add_child(670, n2) 20 | 21 | let all_dfs = t.get_dfs_indices() 22 | for i in range(len(all_dfs)): 23 | print(t[all_dfs[i]]) 24 | 25 | print("---") 26 | 27 | let all_bfs = t.get_bfs_indices() 28 | for i in range(len(all_bfs)): 29 | print(t[all_bfs[i]]) 30 | 31 | 32 | _ = t.swap_nodes(n1, n3) 33 | 34 | print("---") 35 | 36 | let all_bfs2 = t.get_bfs_indices() 37 | for i in range(len(all_bfs2)): 38 | print(t[all_bfs2[i]]) 39 | 40 | 41 | t.print_tree[int_to_str]() 42 | 43 | _ = t.prepend_root(123) 44 | _ = t.prepend_root(321) 45 | t.print_tree[int_to_str]() 46 | 47 | 48 | var t2 = LCRSTreeBuilder(1) 49 | .node(5) 50 | .leaf(7) 51 | .up() 52 | .node(10) 53 | .node(15) 54 | .leaf(13) 55 | .leaf(17) 56 | .up() 57 | .node(45) 58 | .tree() 59 | 60 | t2.print_tree[int_to_str]() 61 | 62 | _ = t.add_tree(t2) 63 | 64 | t.print_tree[int_to_str]() 65 | 66 | t.traverse_dfs[v1]() 67 | 68 | print("-----") 69 | 70 | t.traverse_bfs[v1]() 71 | 72 | let ancestors = t.ancestor_indices(17) 73 | print(t[17], len(ancestors)) 74 | for i in range(len(ancestors)): 75 | print(ancestors[i]) 76 | 77 | 78 | t.compact_dfs() 79 | 80 | print("======") 81 | t.traverse_dfs[v1]() 82 | t.print_tree[int_to_str]() 83 | 84 | t.compact_bfs() 85 | 86 | print("======") 87 | t.traverse_dfs[v1]() 88 | t.print_tree[int_to_str]() 89 | 90 | 91 | print("t len:", t.__len__()) 92 | print("remove node 3") 93 | t.remove(3) 94 | print("t len:", t.__len__()) 95 | t.print_tree[int_to_str]() -------------------------------------------------------------------------------- /trie_dict/dict.mojo: -------------------------------------------------------------------------------- 1 | struct TrieDict[Value: AnyType]: 2 | var chars: DynamicVector[Int8] 3 | var next: DynamicVector[UInt16] 4 | var sibling: DynamicVector[UInt16] 5 | var value_index: DynamicVector[UInt8] 6 | var values: DynamicVector[Value] 7 | var deleted: UInt8 8 | 9 | fn __init__(inout self): 10 | self.chars = DynamicVector[Int8]() 11 | self.next = DynamicVector[UInt16]() 12 | self.sibling = DynamicVector[UInt16]() 13 | self.value_index = DynamicVector[UInt8]() 14 | self.values = DynamicVector[Value]() 15 | self.deleted = 0 16 | 17 | fn __len__(self) -> Int: 18 | return len(self.values) - self.deleted.to_int() 19 | 20 | fn __contains__(self, key: String) -> Bool: 21 | let chars = key._buffer 22 | let key_char_offset = self._find_prefix_count(chars) 23 | let key_offset = key_char_offset.get[0, Int]() 24 | let char_index = key_char_offset.get[1, Int]() 25 | if key_offset == len(chars): 26 | return self.value_index[char_index] > 0 27 | return False 28 | 29 | fn get(self, key: String, default: Value) -> Value: 30 | let chars = key._buffer 31 | let key_char_offset = self._find_prefix_count(chars) 32 | let key_offset = key_char_offset.get[0, Int]() 33 | let char_index = key_char_offset.get[1, Int]() 34 | if key_offset == len(chars): 35 | if self.value_index[char_index] > 0: 36 | return self.values[(self.value_index[char_index] - 1).to_int()] 37 | return default 38 | 39 | fn delete(inout self, key: String): 40 | let chars = key._buffer 41 | let key_char_offset = self._find_prefix_count(chars) 42 | let key_offset = key_char_offset.get[0, Int]() 43 | let char_index = key_char_offset.get[1, Int]() 44 | if key_offset == len(chars): 45 | if self.value_index[char_index] > 0: 46 | self.value_index[char_index] = 0 47 | self.deleted += 1 48 | 49 | 50 | fn put(inout self, key: String, value: Value): 51 | let chars = key._buffer 52 | let key_char_offset = self._find_prefix_count(chars) 53 | let key_offset = key_char_offset.get[0, Int]() 54 | let char_index = key_char_offset.get[1, Int]() 55 | if key_offset == len(chars): 56 | if self.value_index[char_index] == 0: 57 | self.values.push_back(value) 58 | self.value_index[char_index] = UInt8(len(self.values)) 59 | return 60 | # replace 61 | self.values[(self.value_index[char_index] - 1).to_int()] = value 62 | return 63 | if len(self.chars) > char_index: 64 | if char_index == len(self.chars) - 1: 65 | self.next[char_index] = UInt16(len(self.chars)) 66 | else: 67 | self.sibling[char_index] = UInt16(len(self.chars)) 68 | for i in range(key_offset, len(chars)): 69 | self.chars.push_back(chars[i]) 70 | let self_index = UInt16(len(self.sibling)) 71 | self.sibling.push_back(self_index) 72 | self.next.push_back(self_index + 1) 73 | self.value_index.push_back(0) 74 | self.next[len(self.next) - 1] -= 1 75 | self.values.push_back(value) 76 | self.value_index[len(self.next) - 1] = UInt8(len(self.values)) 77 | 78 | fn _find_prefix_count(self, key: DynamicVector[Int8]) -> (Int, Int): 79 | if len(self.chars) == 0: 80 | return 0, 0 81 | var char_index = 0 82 | for key_index in range(len(key)): 83 | while True: 84 | if self.chars[char_index] != key[key_index]: 85 | if self._has_sibling(char_index): 86 | char_index = self.sibling[char_index].to_int() 87 | else: 88 | return key_index, char_index 89 | else: 90 | if self._has_next(char_index): 91 | char_index = self.next[char_index].to_int() 92 | break 93 | else: 94 | return key_index + 1, char_index 95 | return len(key), char_index - 1 96 | 97 | fn _has_sibling(self, char_index: Int) -> Bool: 98 | return self.sibling[char_index].to_int() != char_index 99 | 100 | fn _has_next(self, char_index: Int) -> Bool: 101 | return self.next[char_index].to_int() != char_index 102 | 103 | fn debug(self): 104 | print("Num nodes:", len(self.next)) 105 | print("Num values:", len(self.values)) 106 | var s1: String = "Chars: [" 107 | for i in range(len(self.chars)): 108 | s1 += String(self.chars[i].to_int()) 109 | s1 += "," 110 | s1 += "]" 111 | print(s1) 112 | s1 = "Next: [" 113 | for i in range(len(self.next)): 114 | s1 += String(self.next[i]) 115 | s1 += "," 116 | s1 += "]" 117 | print(s1) 118 | s1 = "Sibling: [" 119 | for i in range(len(self.sibling)): 120 | s1 += String(self.sibling[i]) 121 | s1 += "," 122 | s1 += "]" 123 | print(s1) 124 | s1 = "Value index: [" 125 | for i in range(len(self.value_index)): 126 | s1 += String(self.value_index[i]) 127 | s1 += "," 128 | s1 += "]" 129 | print(s1) 130 | -------------------------------------------------------------------------------- /fiby_tree_int_benchmarks.mojo: -------------------------------------------------------------------------------- 1 | from fiby_tree import FibyTree 2 | from time import now 3 | from random import random_si64 4 | 5 | fn cmp_int(a: Int, b: Int) -> Int: 6 | return a - b 7 | 8 | fn its(a: Int) -> String: 9 | return String(a) 10 | 11 | fn fiby() -> FibyTree[Int, cmp_int, its]: 12 | return FibyTree[Int, cmp_int, its]() 13 | 14 | fn perf_test_random_add(size: Int, min: Int = -30000, max: Int = 30000) -> Float64: 15 | var total = 0 16 | 17 | var tik = now() 18 | var tok = now() 19 | var f = fiby() 20 | 21 | for _ in range(size): 22 | let i = random_si64(min, max).to_int() 23 | tik = now() 24 | f.add(i) 25 | tok = now() 26 | total += (tok - tik) 27 | 28 | return total / size 29 | 30 | fn perf_test_ordered_add(size: Int) -> Float64: 31 | var total = 0 32 | var tik = now() 33 | var f = fiby() 34 | var tok = now() 35 | total += tok - tik 36 | for i in range(size): 37 | tik = now() 38 | f.add(i) 39 | tok = now() 40 | total += (tok - tik) 41 | 42 | total += (tok - tik) 43 | 44 | return total / size 45 | 46 | 47 | fn perf_test_contains(size: Int, balanced: Bool, inout found: Int) -> Float64: 48 | var f = fiby() 49 | for _ in range(size): 50 | let i = random_si64(-size, size).to_int() 51 | f.add(i) 52 | 53 | if balanced: 54 | f.balance() 55 | 56 | var total = 0 57 | 58 | var tik = now() 59 | var tok = now() 60 | 61 | var res = DynamicVector[Bool](size) 62 | for i in range(size): 63 | tik = now() 64 | let r = f.__contains__(i) 65 | tok = now() 66 | res.push_back(r) 67 | total += (tok - tik) 68 | 69 | var count = 0 70 | for i in range(len(res)): 71 | if res[i]: 72 | count += 1 73 | found = count 74 | 75 | return total / size 76 | 77 | fn perf_test_delete(size: Int, balanced: Bool, inout found: Int) -> Float64: 78 | var f = fiby() 79 | for _ in range(size): 80 | let i = random_si64(-size, size).to_int() 81 | f.add(i) 82 | 83 | if balanced: 84 | f.balance() 85 | 86 | var total = 0 87 | 88 | var tik = now() 89 | var tok = now() 90 | 91 | var res = DynamicVector[Bool](size) 92 | for i in range(size): 93 | tik = now() 94 | let r = f.delete(i) 95 | tok = now() 96 | res.push_back(r) 97 | total += (tok - tik) 98 | 99 | var count = 0 100 | for i in range(len(res)): 101 | if res[i]: 102 | count += 1 103 | found = count 104 | 105 | return total / size 106 | 107 | fn perf_test_union(size: Int, balanced: Bool) -> Float64: 108 | var f1 = fiby() 109 | var f2 = fiby() 110 | for _ in range(size): 111 | let i1 = random_si64(-size, size).to_int() 112 | f1.add(i1) 113 | let i2 = random_si64(-size, size).to_int() 114 | f2.add(i2) 115 | if balanced: 116 | f1.balance() 117 | f2.balance() 118 | let s1 = f1.__len__() 119 | let s2 = f2.__len__() 120 | 121 | let tik = now() 122 | f1.union_inplace(f2) 123 | let tok = now() 124 | # print(s1, s2, f1.__len__()) 125 | return (tok - tik) / Float64(size) 126 | 127 | fn perf_test_intersection(size: Int, balanced: Bool) -> Float64: 128 | var f1 = fiby() 129 | var f2 = fiby() 130 | for _ in range(size): 131 | let i1 = random_si64(-size, size).to_int() 132 | f1.add(i1) 133 | let i2 = random_si64(-size, size).to_int() 134 | f2.add(i2) 135 | 136 | if balanced: 137 | f1.balance() 138 | f2.balance() 139 | let s1 = f1.__len__() 140 | let s2 = f2.__len__() 141 | 142 | let tik = now() 143 | f1.intersection_inplace(f2) 144 | let tok = now() 145 | # print(s1, s2, f1.__len__()) 146 | return (tok - tik) / Float64(size) 147 | 148 | fn perf_test_difference(size: Int, balanced: Bool) -> Float64: 149 | var f1 = fiby() 150 | var f2 = fiby() 151 | for _ in range(size): 152 | let i1 = random_si64(-size, size).to_int() 153 | f1.add(i1) 154 | let i2 = random_si64(-size, size).to_int() 155 | f2.add(i2) 156 | if balanced: 157 | f1.balance() 158 | f2.balance() 159 | let s1 = f1.__len__() 160 | let s2 = f2.__len__() 161 | 162 | let tik = now() 163 | f1.difference_inplace(f2) 164 | let tok = now() 165 | # print(s1, s2, f1.__len__()) 166 | return (tok - tik) / Float64(size) 167 | 168 | fn perf_test_symmetric_difference(size: Int, balanced: Bool) -> Float64: 169 | var f1 = fiby() 170 | var f2 = fiby() 171 | for _ in range(size): 172 | let i1 = random_si64(-size, size).to_int() 173 | f1.add(i1) 174 | let i2 = random_si64(-size, size).to_int() 175 | f2.add(i2) 176 | if balanced: 177 | f1.balance() 178 | f2.balance() 179 | let s1 = f1.__len__() 180 | let s2 = f2.__len__() 181 | 182 | let tik = now() 183 | f1.symmetric_difference_inplace(f2) 184 | let tok = now() 185 | 186 | # print(s1, s2, f1.__len__()) 187 | 188 | return (tok - tik) / Float64(size) 189 | 190 | fn main(): 191 | print("===Random Add===") 192 | print(perf_test_random_add(10)) 193 | print(perf_test_random_add(100)) 194 | print(perf_test_random_add(300)) 195 | print(perf_test_random_add(500)) 196 | print(perf_test_random_add(1_000)) 197 | print(perf_test_random_add(3_000)) 198 | print(perf_test_random_add(9_000)) 199 | print(perf_test_random_add(15_000)) 200 | print(perf_test_random_add(30_000)) 201 | print(perf_test_random_add(50_000)) 202 | 203 | print("===Ordered Add===") 204 | 205 | print(perf_test_ordered_add(10)) 206 | print(perf_test_ordered_add(100)) 207 | print(perf_test_ordered_add(300)) 208 | print(perf_test_ordered_add(500)) 209 | print(perf_test_ordered_add(1_000)) 210 | print(perf_test_ordered_add(3_000)) 211 | print(perf_test_ordered_add(9_000)) 212 | print(perf_test_ordered_add(15_000)) 213 | print(perf_test_ordered_add(30_000)) 214 | print(perf_test_ordered_add(50_000)) 215 | 216 | var r = 0 217 | 218 | print("===Contains===") 219 | print(perf_test_contains(10, False, r)) 220 | print(perf_test_contains(100, False, r)) 221 | print(perf_test_contains(300, False, r)) 222 | print(perf_test_contains(500, False, r)) 223 | print(perf_test_contains(1_000, False, r)) 224 | print(perf_test_contains(3_000, False, r)) 225 | print(perf_test_contains(9_000, False, r)) 226 | print(perf_test_contains(15_000, False, r)) 227 | print(perf_test_contains(30_000, False, r)) 228 | print(perf_test_contains(50_000, False, r)) 229 | 230 | print("===Contains Balanced===") 231 | print(perf_test_contains(10, True, r)) 232 | print(perf_test_contains(100, True, r)) 233 | print(perf_test_contains(300, True, r)) 234 | print(perf_test_contains(500, True, r)) 235 | print(perf_test_contains(1_000, True, r)) 236 | print(perf_test_contains(3_000, True, r)) 237 | print(perf_test_contains(9_000, True, r)) 238 | print(perf_test_contains(15_000, True, r)) 239 | print(perf_test_contains(30_000, True, r)) 240 | print(perf_test_contains(50_000, True, r)) 241 | 242 | print("===Delete===") 243 | print(perf_test_delete(10, False, r)) 244 | print(perf_test_delete(100, False, r)) 245 | print(perf_test_delete(300, False, r)) 246 | print(perf_test_delete(500, False, r)) 247 | print(perf_test_delete(1_000, False, r)) 248 | print(perf_test_delete(3_000, False, r)) 249 | print(perf_test_delete(9_000, False, r)) 250 | print(perf_test_delete(15_000, False, r)) 251 | print(perf_test_delete(30_000, False, r)) 252 | print(perf_test_delete(50_000, False, r)) 253 | 254 | print("===Delete Balanced===") 255 | print(perf_test_delete(10, True, r)) 256 | print(perf_test_delete(100, True, r)) 257 | print(perf_test_delete(300, True, r)) 258 | print(perf_test_delete(500, True, r)) 259 | print(perf_test_delete(1_000, True, r)) 260 | print(perf_test_delete(3_000, True, r)) 261 | print(perf_test_delete(9_000, True, r)) 262 | print(perf_test_delete(15_000, True, r)) 263 | print(perf_test_delete(30_000, True, r)) 264 | print(perf_test_delete(50_000, True, r)) 265 | 266 | print("===Union===") 267 | print(perf_test_union(10, False)) 268 | print(perf_test_union(100, False)) 269 | print(perf_test_union(300, False)) 270 | print(perf_test_union(500, False)) 271 | print(perf_test_union(1_000, False)) 272 | print(perf_test_union(3_000, False)) 273 | print(perf_test_union(9_000, False)) 274 | print(perf_test_union(15_000, False)) 275 | print(perf_test_union(30_000, False)) 276 | print(perf_test_union(50_000, False)) 277 | 278 | print("===Union Balanced===") 279 | print(perf_test_union(10, True)) 280 | print(perf_test_union(100, True)) 281 | print(perf_test_union(300, True)) 282 | print(perf_test_union(500, True)) 283 | print(perf_test_union(1_000, True)) 284 | print(perf_test_union(3_000, True)) 285 | print(perf_test_union(9_000, True)) 286 | print(perf_test_union(15_000, True)) 287 | print(perf_test_union(30_000, True)) 288 | print(perf_test_union(50_000, True)) 289 | 290 | print("===Intersection===") 291 | print(perf_test_intersection(10, False)) 292 | print(perf_test_intersection(100, False)) 293 | print(perf_test_intersection(300, False)) 294 | print(perf_test_intersection(500, False)) 295 | print(perf_test_intersection(1_000, False)) 296 | print(perf_test_intersection(3_000, False)) 297 | print(perf_test_intersection(9_000, False)) 298 | print(perf_test_intersection(15_000, False)) 299 | print(perf_test_intersection(30_000, False)) 300 | print(perf_test_intersection(50_000, False)) 301 | 302 | print("===Intersection Balanced===") 303 | print(perf_test_intersection(10, True)) 304 | print(perf_test_intersection(100, True)) 305 | print(perf_test_intersection(300, True)) 306 | print(perf_test_intersection(500, True)) 307 | print(perf_test_intersection(1_000, True)) 308 | print(perf_test_intersection(3_000, True)) 309 | print(perf_test_intersection(9_000, True)) 310 | print(perf_test_intersection(15_000, True)) 311 | print(perf_test_intersection(30_000, True)) 312 | print(perf_test_intersection(50_000, True)) 313 | 314 | print("===Difference===") 315 | print(perf_test_difference(10, False)) 316 | print(perf_test_difference(100, False)) 317 | print(perf_test_difference(300, False)) 318 | print(perf_test_difference(500, False)) 319 | print(perf_test_difference(1_000, False)) 320 | print(perf_test_difference(3_000, False)) 321 | print(perf_test_difference(9_000, False)) 322 | print(perf_test_difference(15_000, False)) 323 | print(perf_test_difference(30_000, False)) 324 | print(perf_test_difference(50_000, False)) 325 | 326 | print("===Difference Balanced===") 327 | print(perf_test_difference(10, True)) 328 | print(perf_test_difference(100, True)) 329 | print(perf_test_difference(300, True)) 330 | print(perf_test_difference(500, True)) 331 | print(perf_test_difference(1_000, True)) 332 | print(perf_test_difference(3_000, True)) 333 | print(perf_test_difference(9_000, True)) 334 | print(perf_test_difference(15_000, True)) 335 | print(perf_test_difference(30_000, True)) 336 | print(perf_test_difference(50_000, True)) 337 | 338 | print("===Symmetric Difference===") 339 | print(perf_test_symmetric_difference(10, False)) 340 | print(perf_test_symmetric_difference(100, False)) 341 | print(perf_test_symmetric_difference(300, False)) 342 | print(perf_test_symmetric_difference(500, False)) 343 | print(perf_test_symmetric_difference(1_000, False)) 344 | print(perf_test_symmetric_difference(3_000, False)) 345 | print(perf_test_symmetric_difference(9_000, False)) 346 | print(perf_test_symmetric_difference(15_000, False)) 347 | print(perf_test_symmetric_difference(30_000, False)) 348 | print(perf_test_symmetric_difference(50_000, False)) 349 | 350 | print("===Symmetric Difference Balanced===") 351 | print(perf_test_symmetric_difference(10, True)) 352 | print(perf_test_symmetric_difference(100, True)) 353 | print(perf_test_symmetric_difference(300, True)) 354 | print(perf_test_symmetric_difference(500, True)) 355 | print(perf_test_symmetric_difference(1_000, True)) 356 | print(perf_test_symmetric_difference(3_000, True)) 357 | print(perf_test_symmetric_difference(9_000, True)) 358 | print(perf_test_symmetric_difference(15_000, True)) 359 | print(perf_test_symmetric_difference(30_000, True)) 360 | print(perf_test_symmetric_difference(50_000, True)) 361 | -------------------------------------------------------------------------------- /fiby_tree_int_tests.mojo: -------------------------------------------------------------------------------- 1 | from testing import * 2 | from fiby_tree import FibyTree 3 | 4 | fn int_eq(a: Int, b: Int) -> Bool: 5 | return a == b 6 | 7 | fn int_cmp(a: Int, b: Int) -> Int: 8 | return a - b 9 | 10 | fn int_to_str(a: Int) -> String: 11 | return String(a) 12 | 13 | fn assert_vec(a: UnsafeFixedVector[Int], b: UnsafeFixedVector[Int]): 14 | if assert_equal(len(a), len(b)): 15 | for i in range(len(a)): 16 | _ = assert_equal(a[i], b[i]) 17 | else: 18 | print("Length", len(a), "is not equal to", len(b)) 19 | 20 | fn vec[T: AnyType](*elements: T) -> UnsafeFixedVector[T]: 21 | let elements_list: VariadicList[T] = elements 22 | var result = UnsafeFixedVector[T](len(elements_list)) 23 | for i in range(len(elements_list)): 24 | result.append(elements[i]) 25 | return result 26 | 27 | fn fiby(*elements: Int) -> FibyTree[Int, int_cmp, int_to_str]: 28 | let elements_list: VariadicList[Int] = elements 29 | var tree = FibyTree[Int, int_cmp, int_to_str]() 30 | for i in range(len(elements_list)): 31 | tree.add(elements[i]) 32 | return tree^ 33 | 34 | fn assert_tree(tree: FibyTree[Int, int_cmp, int_to_str], count: Int, max_depth: Int): 35 | if not assert_equal(tree.__len__(), count): 36 | print("Length assertion failed") 37 | if not assert_equal(tree.max_depth, max_depth): 38 | print("depth assertion failed") 39 | 40 | 41 | 42 | fn test_start_with_empty_tree(): 43 | var bst = fiby() 44 | assert_tree(bst, 0, 0) 45 | bst.add(13) 46 | assert_tree(bst, 1, 1) 47 | bst.add(15) 48 | assert_tree(bst, 2, 2) 49 | _= assert_true(bst.delete(13)) 50 | assert_tree(bst, 1, 2) 51 | _= assert_true(bst.delete(15)) 52 | assert_tree(bst, 0, 2) 53 | _= assert_false(bst.__contains__(15)) 54 | bst.balance() 55 | assert_tree(bst, 0, 0) 56 | 57 | fn test_longer_sequence_dedup_and_balance(): 58 | var bst = fiby(5, 6, 3, 8, 11, 34, 56, 12, 48, 11, 9) 59 | assert_vec(bst.sorted_elements(), vec(3, 5, 6, 8, 9, 11, 12, 34, 48, 56)) 60 | assert_tree(bst, 10, 7) 61 | bst.balance() 62 | assert_tree(bst, 10, 4) 63 | _= assert_true(bst.delete(8)) 64 | assert_vec(bst.sorted_elements(), vec(3, 5, 6, 9, 11, 12, 34, 48, 56)) 65 | assert_tree(bst, 9, 4) 66 | 67 | let elements = bst.sorted_elements() 68 | for i in range(len(elements)): 69 | _= assert_true(bst.delete(elements[i])) 70 | 71 | assert_tree(bst, 0, 4) 72 | bst.add(13) 73 | assert_tree(bst, 1, 4) 74 | bst.clear() 75 | assert_tree(bst, 0, 0) 76 | 77 | 78 | fn test_add_ascending(): 79 | var bst = fiby() 80 | for i in range(10): 81 | bst.add(i) 82 | assert_vec(bst.sorted_elements(), vec(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)) 83 | assert_tree(bst, 10, 10) 84 | bst.balance() 85 | assert_tree(bst, 10, 4) 86 | for i in range(10): 87 | _= assert_true(bst.__contains__(i)) 88 | 89 | fn test_union_inplace(): 90 | var b1 = fiby(1, 2, 3) 91 | b1.union_inplace(fiby()) 92 | assert_vec(b1.sorted_elements(), vec(1, 2, 3)) 93 | 94 | var b2 = fiby() 95 | b2.union_inplace(b1) 96 | assert_vec(b2.sorted_elements(), vec(1, 2, 3)) 97 | 98 | b1.union_inplace(fiby(3, 4, 1)) 99 | assert_vec(b1.sorted_elements(), vec(1, 2, 3, 4)) 100 | 101 | b1.union_inplace(fiby(2, 3)) 102 | assert_vec(b1.sorted_elements(), vec(1, 2, 3, 4)) 103 | 104 | b1.union_inplace(fiby(9, 12, 11, 10)) 105 | assert_vec(b1.sorted_elements(), vec(1, 2, 3, 4, 9, 10, 11, 12)) 106 | 107 | b2 = fiby(1) 108 | b2.union_inplace(fiby(1)) 109 | assert_vec(b2.sorted_elements(), vec(1)) 110 | 111 | b2.union_inplace(fiby(2)) 112 | assert_vec(b2.sorted_elements(), vec(1, 2)) 113 | 114 | b2 = fiby(2) 115 | b2.union_inplace(fiby(1)) 116 | assert_vec(b2.sorted_elements(), vec(1, 2)) 117 | 118 | fn test_union(): 119 | var b1 = fiby(1, 2, 3) 120 | var b3 = b1.union(fiby()) 121 | assert_vec(b3.sorted_elements(), vec(1, 2, 3)) 122 | 123 | var b2 = fiby() 124 | b3 = b2.union(b1) 125 | assert_vec(b3.sorted_elements(), vec(1, 2, 3)) 126 | 127 | b3 = b1.union(fiby(3, 4, 1)) 128 | assert_vec(b3.sorted_elements(), vec(1, 2, 3, 4)) 129 | 130 | b3 = b1.union(fiby(2, 3)) 131 | assert_vec(b3.sorted_elements(), vec(1, 2, 3, 4)) 132 | 133 | b3 = b1.union(fiby(9, 12, 11, 10)) 134 | assert_vec(b3.sorted_elements(), vec(1, 2, 3, 4, 9, 10, 11, 12)) 135 | 136 | b2 = fiby(1) 137 | b3 = b2.union(fiby(1)) 138 | assert_vec(b3.sorted_elements(), vec(1)) 139 | 140 | b2.union_inplace(fiby(2)) 141 | assert_vec(b2.sorted_elements(), vec(1, 2)) 142 | 143 | b2 = fiby(2) 144 | b2.union_inplace(fiby(1)) 145 | assert_vec(b2.sorted_elements(), vec(1, 2)) 146 | 147 | 148 | fn test_intersection_inplace(): 149 | var b1 = fiby(1, 2, 3) 150 | b1.intersection_inplace(fiby(3, 4, 1, 6, 7, 10)) 151 | assert_vec(b1.sorted_elements(), vec(1, 3)) 152 | 153 | b1.intersection_inplace(fiby()) 154 | assert_vec(b1.sorted_elements(), vec[Int]()) 155 | 156 | b1.intersection_inplace(fiby(3, 4, 1)) 157 | assert_vec(b1.sorted_elements(), vec[Int]()) 158 | 159 | var b2 = fiby(3, 4, 1, 6, 7, 10) 160 | b2.intersection_inplace(fiby(1, 2, 3, 8)) 161 | assert_vec(b2.sorted_elements(), vec(1, 3)) 162 | 163 | b2 = fiby(1) 164 | b2.intersection_inplace(fiby(1)) 165 | assert_vec(b2.sorted_elements(), vec(1)) 166 | 167 | b2 = fiby(1) 168 | b2.intersection_inplace(fiby(2)) 169 | assert_vec(b2.sorted_elements(), vec[Int]()) 170 | 171 | b2 = fiby(2) 172 | b2.intersection_inplace(fiby(1)) 173 | assert_vec(b2.sorted_elements(), vec[Int]()) 174 | 175 | fn test_difference_inplace(): 176 | var b1 = fiby(1, 2, 3) 177 | b1.difference_inplace(fiby(5, 6, 7, 1)) 178 | assert_vec(b1.sorted_elements(), vec(2, 3)) 179 | b1.difference_inplace(fiby()) 180 | assert_vec(b1.sorted_elements(), vec(2, 3)) 181 | b1.difference_inplace(fiby(1, 12, 34)) 182 | assert_vec(b1.sorted_elements(), vec(2, 3)) 183 | 184 | var b2 = fiby() 185 | b2.difference_inplace(fiby(1, 2, 3)) 186 | assert_tree(b2, 0, 0) 187 | 188 | b2 = fiby(1) 189 | b2.difference_inplace(fiby(1)) 190 | assert_vec(b2.sorted_elements(), vec[Int]()) 191 | 192 | b2 = fiby(1) 193 | b2.difference_inplace(fiby(2)) 194 | assert_vec(b2.sorted_elements(), vec(1)) 195 | 196 | b2 = fiby(2) 197 | b2.difference_inplace(fiby(1)) 198 | assert_vec(b2.sorted_elements(), vec(2)) 199 | 200 | fn test_other_difference_inplace(): 201 | var b1 = fiby(1, 2, 3) 202 | b1.other_difference_inplace(fiby(5, 6, 7, 1)) 203 | assert_vec(b1.sorted_elements(), vec(5, 6, 7)) 204 | b1.other_difference_inplace(fiby()) 205 | assert_vec(b1.sorted_elements(), vec[Int]()) 206 | 207 | b1 = fiby(1, 2, 3) 208 | b1.other_difference_inplace(fiby(0, 1, 12, 34)) 209 | assert_vec(b1.sorted_elements(), vec(0, 12, 34)) 210 | 211 | var b2 = fiby() 212 | b2.other_difference_inplace(fiby(1, 2, 3)) 213 | assert_vec(b2.sorted_elements(), vec(1, 2, 3)) 214 | 215 | b2 = fiby(1) 216 | b2.other_difference_inplace(fiby(1)) 217 | assert_vec(b2.sorted_elements(), vec[Int]()) 218 | 219 | b2 = fiby(1) 220 | b2.other_difference_inplace(fiby(2)) 221 | assert_vec(b2.sorted_elements(), vec(2)) 222 | 223 | b2 = fiby(2) 224 | b2.other_difference_inplace(fiby(1)) 225 | assert_vec(b2.sorted_elements(), vec(1)) 226 | 227 | fn test_symmetric_difference_inplace(): 228 | var b1 = fiby(1, 2, 3) 229 | b1.symmetric_difference_inplace(fiby(3, 4, 5)) 230 | assert_vec(b1.sorted_elements(), vec(1, 2, 4, 5)) 231 | 232 | b1.symmetric_difference_inplace(fiby(0, 2, 8, 5, 13)) 233 | assert_vec(b1.sorted_elements(), vec(0, 1, 4, 8, 13)) 234 | 235 | b1.symmetric_difference_inplace(fiby()) 236 | assert_vec(b1.sorted_elements(), vec(0, 1, 4, 8, 13)) 237 | 238 | var b2 = fiby() 239 | b2.symmetric_difference_inplace(fiby(1, 2, 3)) 240 | assert_vec(b2.sorted_elements(), vec(1, 2, 3)) 241 | 242 | b2 = fiby() 243 | b2.symmetric_difference_inplace(fiby(1)) 244 | assert_vec(b2.sorted_elements(), vec(1)) 245 | 246 | b2 = fiby(1) 247 | b2.symmetric_difference_inplace(fiby()) 248 | assert_vec(b2.sorted_elements(), vec(1)) 249 | 250 | b2 = fiby(1) 251 | b2.symmetric_difference_inplace(fiby(1)) 252 | assert_vec(b2.sorted_elements(), vec[Int]()) 253 | 254 | b2 = fiby(1) 255 | b2.symmetric_difference_inplace(fiby(2)) 256 | assert_vec(b2.sorted_elements(), vec(1, 2)) 257 | 258 | b2 = fiby(2) 259 | b2.symmetric_difference_inplace(fiby(1)) 260 | assert_vec(b2.sorted_elements(), vec(1, 2)) 261 | 262 | fn test_disjoint(): 263 | _= assert_true(fiby().is_disjoint(fiby())) 264 | _= assert_true(fiby().is_disjoint(fiby(1))) 265 | _= assert_true(fiby(1).is_disjoint(fiby())) 266 | _= assert_true(fiby(1).is_disjoint(fiby(2))) 267 | _= assert_false(fiby(1).is_disjoint(fiby(1))) 268 | _= assert_true(fiby(1, 3).is_disjoint(fiby(2, 5, 6))) 269 | _= assert_true(fiby(1, 3, 5).is_disjoint(fiby(2, 0, 7))) 270 | _= assert_false(fiby(1, 3, 5).is_disjoint(fiby(2, 5))) 271 | _= assert_false(fiby(1, 5).is_disjoint(fiby(2, 5))) 272 | 273 | fn test_subset(): 274 | _= assert_true(fiby().is_subset(fiby())) 275 | _= assert_true(fiby().is_subset(fiby(1, 2, 3))) 276 | _= assert_true(fiby(3).is_subset(fiby(3))) 277 | _= assert_true(fiby(3).is_subset(fiby(1, 2, 3))) 278 | _= assert_true(fiby(3, 1).is_subset(fiby(1, 2, 3))) 279 | _= assert_true(fiby(3, 1, 2).is_subset(fiby(1, 2, 3))) 280 | _= assert_false(fiby(1).is_subset(fiby(3))) 281 | _= assert_false(fiby(3, 1, 2, 5).is_subset(fiby(1, 2, 3))) 282 | _= assert_false(fiby(3, 1, 5).is_subset(fiby(1, 2, 3))) 283 | 284 | fn test_superset(): 285 | _= assert_false(fiby(1).is_superset(fiby(2))) 286 | _= assert_false(fiby(1, 5, 8).is_superset(fiby(1, 5, 8, 9))) 287 | _= assert_false(fiby(1, 5, 8).is_superset(fiby(1, 5, 9))) 288 | 289 | _= assert_true(fiby().is_superset(fiby())) 290 | _= assert_true(fiby(1).is_superset(fiby(1))) 291 | _= assert_true(fiby(1, 5, 8, 9).is_superset(fiby(1, 5, 8, 9))) 292 | _= assert_true(fiby(1, 5, 8, 9).is_superset(fiby(1, 5, 8))) 293 | _= assert_true(fiby(0, 1, 5, 8, 9).is_superset(fiby(1, 5, 8))) 294 | 295 | fn test_min_index(): 296 | _= assert_equal(fiby().min_index(), -1) 297 | _= assert_equal(fiby(1).min_index(), 0) 298 | _= assert_equal(fiby(1, 2, 3, 4).min_index(), 0) 299 | _= assert_equal(fiby(3, 4, 1, 2).min_index(), 2) 300 | var f = fiby(3, 4, 1, 2) 301 | f.balance() 302 | _= assert_equal(f.min_index(), 3) 303 | 304 | f = fiby(1, 3, 2, 4, 5) 305 | _= assert_equal(f.min_index(), 0) 306 | f = fiby(3, 1, 2, 4, 5) 307 | _= assert_equal(f.min_index(), 1) 308 | f = fiby(3, 2, 1, 4, 5) 309 | _= assert_equal(f.min_index(), 2) 310 | f = fiby(3, 2, 4, 1, 5) 311 | _= assert_equal(f.min_index(), 3) 312 | f = fiby(3, 2, 4, 5, 1) 313 | _= assert_equal(f.min_index(), 4) 314 | 315 | f = fiby(1, 3, 2, 4, 5) 316 | f.balance() 317 | _= assert_equal(f.min_index(), 3) 318 | f = fiby(3, 1, 2, 4, 5) 319 | f.balance() 320 | _= assert_equal(f.min_index(), 3) 321 | f = fiby(3, 2, 1, 4, 5) 322 | f.balance() 323 | _= assert_equal(f.min_index(), 3) 324 | f = fiby(3, 2, 4, 1, 5) 325 | f.balance() 326 | _= assert_equal(f.min_index(), 3) 327 | f = fiby(3, 2, 4, 5, 1) 328 | f.balance() 329 | _= assert_equal(f.min_index(), 3) 330 | 331 | fn test_max_index(): 332 | _= assert_equal(fiby().max_index(), -1) 333 | _= assert_equal(fiby(1).max_index(), 0) 334 | _= assert_equal(fiby(1, 2, 3, 4).max_index(), 3) 335 | _= assert_equal(fiby(3, 4, 1, 2).max_index(), 1) 336 | var f = fiby(3, 4, 1, 2) 337 | f.balance() 338 | _= assert_equal(f.max_index(), 2) 339 | 340 | f = fiby(3, 4, 1, 2, 5, 6, 0) 341 | f.balance() 342 | _= assert_equal(f.max_index(), 6) 343 | _= assert_equal(f.elements[6], 6) 344 | 345 | f = fiby(1, 3, 2, 4, 5) 346 | _= assert_equal(f.max_index(), 4) 347 | f = fiby(3, 1, 2, 5, 4) 348 | _= assert_equal(f.max_index(), 3) 349 | f = fiby(3, 2, 5, 4, 1) 350 | _= assert_equal(f.max_index(), 2) 351 | f = fiby(3, 5, 4, 1, 2) 352 | _= assert_equal(f.max_index(), 1) 353 | f = fiby(5, 2, 4, 5, 3) 354 | _= assert_equal(f.max_index(), 0) 355 | 356 | f = fiby(1, 3, 2, 4, 5) 357 | f.balance() 358 | _= assert_equal(f.max_index(), 2) 359 | f = fiby(3, 1, 2, 5, 4) 360 | f.balance() 361 | _= assert_equal(f.max_index(), 2) 362 | f = fiby(3, 2, 5, 4, 1) 363 | f.balance() 364 | _= assert_equal(f.max_index(), 2) 365 | f = fiby(3, 5, 4, 1, 2) 366 | f.balance() 367 | _= assert_equal(f.max_index(), 2) 368 | f = fiby(5, 2, 4, 5, 3) 369 | f.balance() 370 | _= assert_equal(f.max_index(), 2) 371 | _= assert_equal(f.elements[2], 5) 372 | 373 | fn main(): 374 | test_start_with_empty_tree() 375 | test_longer_sequence_dedup_and_balance() 376 | test_add_ascending() 377 | test_union_inplace() 378 | test_intersection_inplace() 379 | test_difference_inplace() 380 | test_other_difference_inplace() 381 | test_symmetric_difference_inplace() 382 | # Uncomment once https://github.com/modularml/mojo/issues/500 is shipped 383 | # test_union() 384 | test_disjoint() 385 | test_subset() 386 | test_superset() 387 | test_min_index() 388 | test_max_index() 389 | 390 | print("SUCCESS!!!") -------------------------------------------------------------------------------- /left_child_right_sibling/lcrs_tree.mojo: -------------------------------------------------------------------------------- 1 | struct LCRSTree[T: AnyType]: 2 | var elements: DynamicVector[T] 3 | var left_child: DynamicVector[UInt16] 4 | var right_sibling: DynamicVector[UInt16] 5 | var parent: DynamicVector[UInt16] 6 | var deleted: DynamicVector[UInt16] 7 | 8 | fn __init__(inout self, root: T): 9 | self.elements = DynamicVector[T]() 10 | self.left_child = DynamicVector[UInt16]() 11 | self.right_sibling = DynamicVector[UInt16]() 12 | self.parent = DynamicVector[UInt16]() 13 | self.deleted = DynamicVector[UInt16]() 14 | self.elements.push_back(root) 15 | self.left_child.push_back(0) 16 | self.right_sibling.push_back(0) 17 | self.parent.push_back(0) 18 | 19 | fn __moveinit__(inout self, owned existing: Self): 20 | self.elements = existing.elements 21 | self.left_child = existing.left_child 22 | self.right_sibling = existing.right_sibling 23 | self.parent = existing.parent 24 | self.deleted = existing.deleted 25 | 26 | fn add_child(inout self, node: T, parent_index: UInt16 = 0) -> UInt16: 27 | var node_index = len(self.elements) 28 | if len(self.deleted) == 0: 29 | self.elements.push_back(node) 30 | self.left_child.push_back(node_index) 31 | self.right_sibling.push_back(node_index) 32 | self.parent.push_back(parent_index) 33 | else: 34 | node_index = self.deleted.pop_back().to_int() 35 | self.elements[node_index] = node 36 | self.left_child[node_index] = node_index 37 | self.right_sibling[node_index] = node_index 38 | self.parent[node_index] = parent_index 39 | 40 | self._add_as_last_child(parent_index, node_index) 41 | return node_index 42 | 43 | fn add_tree(inout self, tree: Self, parent_index: UInt16 = 0) -> UInt16: 44 | let new_root = len(self.elements) 45 | for i in range(len(tree.elements)): 46 | self.elements.push_back(tree.elements[i]) 47 | self.left_child.push_back(tree.left_child[i] + new_root) 48 | self.right_sibling.push_back(tree.right_sibling[i] + new_root) 49 | self.parent.push_back(tree.parent[i] + new_root) 50 | 51 | self.parent[new_root] = 0 52 | 53 | for i in range(len(tree.deleted)): 54 | self.deleted.push_back(tree.deleted[i] + new_root) 55 | 56 | self._add_as_last_child(parent_index, new_root) 57 | return new_root 58 | 59 | fn _add_as_last_child(inout self, parent_index: UInt16, node_index: UInt16): 60 | var current_child = self.left_child[parent_index.to_int()] 61 | if current_child == parent_index: 62 | self.left_child[parent_index.to_int()] = node_index 63 | else: 64 | var right_sibling = self.right_sibling[current_child.to_int()] 65 | while right_sibling != current_child: 66 | current_child = right_sibling 67 | right_sibling = self.right_sibling[current_child.to_int()] 68 | self.right_sibling[current_child.to_int()] = node_index 69 | 70 | fn prepend_root(inout self, element: T) -> UInt16: 71 | var old_parent_index = len(self.elements) 72 | 73 | if len(self.deleted) == 0: 74 | self.elements.push_back(self.elements[0]) 75 | self.left_child.push_back(self.left_child[0]) 76 | self.right_sibling.push_back(old_parent_index) 77 | self.parent.push_back(0) 78 | else: 79 | old_parent_index = self.deleted.pop_back().to_int() 80 | self.elements[old_parent_index] = self.elements[0] 81 | self.left_child[old_parent_index] = self.left_child[0] 82 | self.right_sibling[old_parent_index] = old_parent_index 83 | self.parent[old_parent_index] = 0 84 | 85 | self.elements[0] = element 86 | self.left_child[0] = old_parent_index 87 | 88 | var child = self.left_child[old_parent_index] 89 | self.parent[child.to_int()] = old_parent_index 90 | while self.right_sibling[child.to_int()] != child: 91 | child = self.right_sibling[child.to_int()] 92 | self.parent[child.to_int()] = old_parent_index 93 | 94 | return old_parent_index 95 | 96 | fn __getitem__(self, node_index: UInt16) -> T: 97 | return self.elements[node_index.to_int()] 98 | 99 | fn __setitem__(inout self, node_index: UInt16, element: T): 100 | self.elements[node_index.to_int()] = element 101 | 102 | fn __len__(self: Self) -> Int: 103 | return len(self.elements) - len(self.deleted) 104 | 105 | fn swap_elements(inout self, index_a: UInt16, index_b: UInt16): 106 | let temp = self.elements[index_a.to_int()] 107 | self.elements[index_a.to_int()] = self.elements[index_b.to_int()] 108 | self.elements[index_b.to_int()] = temp 109 | 110 | fn swap_nodes(inout self, index_a: UInt16, index_b: UInt16) -> Bool: 111 | if index_a == index_b: 112 | return False 113 | if self.is_leaf(index_a) and self.is_leaf(index_b): 114 | self.swap_elements(index_a, index_b) 115 | return True 116 | if self.are_siblings(index_a, index_b): 117 | self._swap_siblings(index_a, index_b) 118 | return True 119 | if not self.is_root(index_a) and not self.is_root(index_b): 120 | self._swap_nodes(index_a, index_b) 121 | return False 122 | 123 | fn is_leaf(self, index: UInt16) -> Bool: 124 | return (self.left_child[index.to_int()] == index).__bool__() 125 | 126 | fn is_root(self, index: UInt16) -> Bool: 127 | return (self.parent[index.to_int()] == index).__bool__() 128 | 129 | fn are_siblings(self, index_a: UInt16, index_b: UInt16) -> Bool: 130 | return (self.parent[index_a.to_int()] == self.parent[index_b.to_int()]).__bool__() 131 | 132 | fn _swap_siblings(inout self, index_a: UInt16, index_b: UInt16): 133 | let parent = self.parent[index_a.to_int()] 134 | 135 | let children = self.children_indices(parent) 136 | var children_index_left = -1 137 | var children_index_right = -1 138 | for i in range(len(children)): 139 | if children[i] == index_a or children[i] == index_b: 140 | if children_index_left == -1: 141 | children_index_left = i 142 | else: 143 | children_index_right = i 144 | break 145 | 146 | let left_sibling = children[children_index_left] 147 | let right_sibling = children[children_index_right] 148 | let left_sibling_of_right = children[children_index_right - 1] 149 | 150 | if children_index_left == 0: 151 | self.left_child[parent.to_int()] = right_sibling 152 | 153 | let prev_right_sibling_of_left = self.right_sibling[left_sibling.to_int()] 154 | if children_index_right == len(children) - 1: 155 | self.right_sibling[left_sibling.to_int()] = left_sibling 156 | else: 157 | self.right_sibling[left_sibling.to_int()] = self.right_sibling[right_sibling.to_int()] 158 | 159 | if left_sibling == left_sibling_of_right: 160 | self.right_sibling[right_sibling.to_int()] = left_sibling 161 | else: 162 | self.right_sibling[right_sibling.to_int()] = prev_right_sibling_of_left 163 | self.right_sibling[left_sibling_of_right.to_int()] = left_sibling 164 | 165 | fn _swap_nodes(inout self, index_a: UInt16, index_b: UInt16): 166 | let parent_a = self.parent[index_a.to_int()] 167 | let parent_b = self.parent[index_b.to_int()] 168 | 169 | self.parent[index_a.to_int()] = parent_b 170 | self.parent[index_b.to_int()] = parent_a 171 | 172 | let right_sibling_a = self.right_sibling[index_a.to_int()] 173 | let right_sibling_b = self.right_sibling[index_b.to_int()] 174 | 175 | if self.left_child[parent_a.to_int()] == index_a: 176 | self.left_child[parent_a.to_int()] = index_b 177 | else: 178 | var left_sibling = self.left_child[parent_a.to_int()] 179 | while self.right_sibling[left_sibling.to_int()] != index_a: 180 | left_sibling = self.right_sibling[left_sibling.to_int()] 181 | self.right_sibling[left_sibling.to_int()] = index_b 182 | 183 | if self.left_child[parent_b.to_int()] == index_b: 184 | self.left_child[parent_b.to_int()] = index_a 185 | else: 186 | var left_sibling = self.left_child[parent_b.to_int()] 187 | while self.right_sibling[left_sibling.to_int()] != index_b: 188 | left_sibling = self.right_sibling[left_sibling.to_int()] 189 | self.right_sibling[left_sibling.to_int()] = index_a 190 | 191 | if right_sibling_a != index_a: 192 | self.right_sibling[index_b.to_int()] = right_sibling_a 193 | else: 194 | self.right_sibling[index_b.to_int()] = index_b 195 | 196 | if right_sibling_b != index_b: 197 | self.right_sibling[index_a.to_int()] = right_sibling_b 198 | else: 199 | self.right_sibling[index_a.to_int()] = index_a 200 | 201 | fn children_indices(self, parent_index: UInt16) -> DynamicVector[UInt16]: 202 | var result = DynamicVector[UInt16]() 203 | var current_child = self.left_child[parent_index.to_int()] 204 | if current_child == parent_index: 205 | return result 206 | result.push_back(current_child) 207 | var right_sibling = self.right_sibling[current_child.to_int()] 208 | while right_sibling != current_child: 209 | result.push_back(right_sibling) 210 | current_child = right_sibling 211 | right_sibling = self.right_sibling[current_child.to_int()] 212 | return result 213 | 214 | fn children_count(self, parent_index: UInt16) -> Int: 215 | var result = 0 216 | var current_child = self.left_child[parent_index.to_int()] 217 | if current_child == parent_index: 218 | return result 219 | result += 1 220 | var right_sibling = self.right_sibling[current_child.to_int()] 221 | while right_sibling != current_child: 222 | result += 1 223 | current_child = right_sibling 224 | right_sibling = self.right_sibling[current_child.to_int()] 225 | return result 226 | 227 | fn ancestor_indices(self, child_index: UInt16) -> DynamicVector[UInt16]: 228 | var result = DynamicVector[UInt16]() 229 | var ancestor = child_index 230 | while ancestor != self.parent[ancestor.to_int()]: 231 | ancestor = self.parent[ancestor.to_int()] 232 | result.push_back(ancestor) 233 | return result 234 | 235 | fn get_dfs_indices(self, root: UInt16 = 0) -> DynamicVector[UInt16]: 236 | var result = DynamicVector[UInt16]() 237 | if len(self.elements) <= root.to_int(): 238 | return result 239 | result.push_back(root) 240 | self._dfs(root, result) 241 | return result 242 | 243 | fn _dfs(self, index: UInt16, inout result: DynamicVector[UInt16]): 244 | let child = self.left_child[index.to_int()] 245 | if child == index: 246 | let sibling = self.right_sibling[index.to_int()] 247 | if sibling == index: 248 | return 249 | result.push_back(sibling) 250 | self._dfs(sibling, result) 251 | else: 252 | result.push_back(child) 253 | self._dfs(child, result) 254 | let sibling = self.right_sibling[index.to_int()] 255 | if sibling == index: 256 | return 257 | result.push_back(sibling) 258 | self._dfs(sibling, result) 259 | 260 | fn traverse_dfs[visitor: fn(UInt16, T) -> Bool](self, root: UInt16 = 0) -> None: 261 | if len(self.elements) > root.to_int(): 262 | self._traverse_dfs[visitor](root, True) 263 | 264 | fn _traverse_dfs[visitor: fn(UInt16, T) -> Bool](self, index: UInt16, root: Bool = False) -> None: 265 | if visitor(index, self.elements[index.to_int()]): 266 | let child = self.left_child[index.to_int()] 267 | if child == index: 268 | if root: 269 | return 270 | let sibling = self.right_sibling[index.to_int()] 271 | if sibling == index: 272 | return 273 | self._traverse_dfs[visitor](sibling) 274 | else: 275 | self._traverse_dfs[visitor](child) 276 | let sibling = self.right_sibling[index.to_int()] 277 | if sibling == index: 278 | return 279 | self._traverse_dfs[visitor](sibling) 280 | 281 | fn get_bfs_indices(self, root: UInt16 = 0) -> DynamicVector[UInt16]: 282 | var result = DynamicVector[UInt16]() 283 | if len(self.elements) <= root.to_int(): 284 | return result 285 | result.push_back(root) 286 | var visited = 0 287 | while len(result) > visited: 288 | let index = result[visited] 289 | var child = self.left_child[result[visited].to_int()] 290 | if child != index: 291 | result.push_back(child) 292 | var sibling = self.right_sibling[child.to_int()] 293 | while sibling != child: 294 | result.push_back(sibling) 295 | child = sibling 296 | sibling = self.right_sibling[child.to_int()] 297 | visited += 1 298 | return result 299 | 300 | fn traverse_bfs[visitor: fn(UInt16, T) -> Bool](self, root: UInt16 = 0) -> None: 301 | if len(self.elements) <= root.to_int(): 302 | return 303 | if visitor(root, self.elements[root.to_int()]): 304 | var visited = 0 305 | # Could be implemented with stack allocated pointer 306 | var result = UnsafeFixedVector[UInt16](len(self.elements)) 307 | result.append(root) 308 | while len(result) > visited: 309 | let index = result[visited] 310 | var child = self.left_child[result[visited].to_int()] 311 | if child != index: 312 | result.append(child) 313 | if not visitor(child, self.elements[child.to_int()]): 314 | return 315 | var sibling = self.right_sibling[child.to_int()] 316 | while sibling != child: 317 | result.append(sibling) 318 | if not visitor(sibling, self.elements[sibling.to_int()]): 319 | return 320 | child = sibling 321 | sibling = self.right_sibling[child.to_int()] 322 | visited += 1 323 | 324 | fn remove(inout self, index: UInt16): 325 | if index == 0: 326 | self.elements.clear() 327 | self.left_child.clear() 328 | self.right_sibling.clear() 329 | self.parent.clear() 330 | return 331 | let parent = self.parent[index.to_int()] 332 | if parent != index: 333 | var parent_child = self.left_child[parent.to_int()] 334 | if parent_child == index: 335 | if self.right_sibling[index.to_int()] != index: 336 | self.left_child[parent.to_int()] = self.right_sibling[index.to_int()] 337 | else: 338 | self.left_child[parent.to_int()] = parent 339 | else: 340 | while self.right_sibling[parent_child.to_int()] != index: 341 | parent_child = self.right_sibling[parent_child.to_int()] 342 | if self.right_sibling[index.to_int()] == index: 343 | self.right_sibling[parent_child.to_int()] = parent_child 344 | else: 345 | self.right_sibling[parent_child.to_int()] = self.right_sibling[index.to_int()] 346 | let indecies_to_remove = self.get_bfs_indices(index) 347 | for i in range(len(indecies_to_remove)): 348 | self.deleted.push_back(indecies_to_remove[i]) 349 | 350 | fn compact_dfs(inout self, root_index: UInt16 = 0): 351 | let indices = self.get_dfs_indices(root_index) 352 | self._compact(indices) 353 | 354 | fn compact_bfs(inout self, root_index: UInt16 = 0): 355 | let indices = self.get_bfs_indices(root_index) 356 | self._compact(indices) 357 | 358 | fn _compact(inout self, indices: DynamicVector[UInt16]): 359 | let old_len = len(self.elements) 360 | let new_len = len(indices) 361 | var map = DynamicVector[Int](old_len) 362 | for i in range(new_len): 363 | map[indices[i].to_int()] = i 364 | var new_elements = DynamicVector[T](new_len) 365 | for i in range(new_len): 366 | new_elements.push_back(self.elements[indices[i].to_int()]) 367 | self.elements = new_elements 368 | var new_left_child = DynamicVector[UInt16](new_len) 369 | for i in range(new_len): 370 | new_left_child.push_back(map[self.left_child[indices[i].to_int()].to_int()]) 371 | self.left_child = new_left_child 372 | var new_right_sibling = DynamicVector[UInt16](new_len) 373 | for i in range(new_len): 374 | new_right_sibling.push_back(map[self.right_sibling[indices[i].to_int()].to_int()]) 375 | self.right_sibling = new_right_sibling 376 | var new_parent = DynamicVector[UInt16](new_len) 377 | for i in range(new_len): 378 | new_parent.push_back(map[self.parent[indices[i].to_int()].to_int()]) 379 | self.parent = new_parent 380 | self.deleted.clear() 381 | 382 | fn print_tree[to_str: fn(T) -> String](inout self, root: UInt16 = 0): 383 | self._print[to_str]("", 0) 384 | 385 | fn _print[to_str: fn(T) -> String](inout self, indentation: String, index: UInt16): 386 | if len(indentation) > 10: 387 | return 388 | if len(indentation) > 0: 389 | print(indentation, "-", to_str(self[index.to_int()])) 390 | else: 391 | print("-", to_str(self[index.to_int()])) 392 | 393 | let children = self.children_indices(index) 394 | for i in range(len(children)): 395 | self._print[to_str](indentation + " ", children[i]) -------------------------------------------------------------------------------- /fiby_tree/fiby_tree.mojo: -------------------------------------------------------------------------------- 1 | from math.bit import bit_length 2 | from utils.list import VariadicList 3 | 4 | struct FibyTree[T: AnyType, cmp: fn(T, T) -> Int, to_str: fn(T) -> String]: 5 | alias Union = 0 6 | alias Intersection = 1 7 | alias Difference = 2 8 | alias SymetricDifference = 3 9 | alias OtherDifference = 4 10 | alias IsDisjoint = 5 11 | alias IsSubset = 6 12 | alias IsSuperset = 7 13 | 14 | var elements: DynamicVector[T] 15 | var left: DynamicVector[UInt16] 16 | var right: DynamicVector[UInt16] 17 | var deleted: Int 18 | var max_depth: UInt16 19 | var balanced: Bool 20 | 21 | fn __init__(inout self, *elements: T): 22 | self.elements = DynamicVector[T]() 23 | self.left = DynamicVector[UInt16]() 24 | self.right = DynamicVector[UInt16]() 25 | self.deleted = 0 26 | self.max_depth = 0 27 | self.balanced = False 28 | 29 | let elements_list: VariadicList[T] = elements 30 | for i in range(len(elements_list)): 31 | self.add(elements[i]) 32 | 33 | fn __moveinit__(inout self, owned existing: Self): 34 | self.elements = existing.elements 35 | self.left = existing.left 36 | self.right = existing.right 37 | self.deleted = existing.deleted 38 | self.max_depth = existing.max_depth 39 | self.balanced = existing.balanced 40 | 41 | @always_inline("nodebug") 42 | fn has_left(self, parent: UInt16) -> Bool: 43 | return (self.left[parent.to_int()] != parent).__bool__() 44 | 45 | @always_inline("nodebug") 46 | fn has_right(self, parent: UInt16) -> Bool: 47 | return (self.right[parent.to_int()] != parent).__bool__() 48 | 49 | fn add(inout self, element: T): 50 | if self.__len__() == 0: 51 | self._set_root(element) 52 | self._set_max_depth(1) 53 | return 54 | var parent = 0 55 | var depth: UInt16 = 1 56 | while True: 57 | let diff = cmp(self.elements[parent], element) 58 | if diff == 0: 59 | return 60 | depth += 1 61 | if diff > 0: 62 | let left = self.left[parent].to_int() 63 | if left == parent: 64 | self._add_left(parent, element) 65 | break 66 | else: 67 | parent = left 68 | else: 69 | let right = self.right[parent].to_int() 70 | if right == parent: 71 | self._add_right(parent, element) 72 | break 73 | else: 74 | parent = right 75 | 76 | self.balanced = False 77 | self._set_max_depth(depth) 78 | if self.max_depth > self._optimal_depth() ** 2: 79 | self.balance() 80 | 81 | @always_inline("nodebug") 82 | fn _set_max_depth(inout self, candidate: UInt16): 83 | if self.max_depth < candidate: 84 | self.max_depth = candidate 85 | 86 | fn _optimal_depth(self) -> UInt16: 87 | return bit_length(UInt16(self.__len__())) 88 | 89 | @always_inline("nodebug") 90 | fn _set_root(inout self, element: T): 91 | if len(self.elements) == 0: 92 | self.elements.push_back(element) 93 | self.left.push_back(0) 94 | self.right.push_back(0) 95 | else: 96 | self.elements[0] = element 97 | self.left[0] = 0 98 | self.right[0] = 0 99 | if self.deleted > 0: 100 | self.deleted -= 1 101 | 102 | @always_inline("nodebug") 103 | fn _add_left(inout self, parent: UInt16, element: T): 104 | let index = len(self.elements) 105 | self.elements.push_back(element) 106 | self.left.push_back(index) 107 | self.right.push_back(index) 108 | self.left[parent.to_int()] = index 109 | 110 | @always_inline("nodebug") 111 | fn _add_right(inout self, parent: UInt16, element: T): 112 | let index = len(self.elements) 113 | self.elements.push_back(element) 114 | self.left.push_back(index) 115 | self.right.push_back(index) 116 | self.right[parent.to_int()] = index 117 | 118 | @always_inline("nodebug") 119 | fn delete(inout self, element: T) -> Bool: 120 | let index_tuple = self._get_index(element) 121 | let parent = index_tuple.get[0, Int]() 122 | let index = index_tuple.get[1, Int]() 123 | if index == -1: 124 | return False 125 | 126 | self.balanced = False 127 | 128 | if self._is_leaf(index): 129 | self._delete_leaf(index, parent) 130 | return True 131 | 132 | if self.has_left(index) and not self.has_right(index): 133 | if index == 0: 134 | let left = self.left[0] 135 | self.elements[0] = self.elements[left.to_int()] 136 | if self.has_left(left): 137 | self.left[0] = self.left[left.to_int()] 138 | else: 139 | self.left[0] = 0 140 | if self.has_right(left): 141 | self.right[0] = self.right[left.to_int()] 142 | else: 143 | self.right[0] = 0 144 | else: 145 | if self.left[parent] == index: 146 | self.left[parent] = self.left[index] 147 | else: 148 | self.right[parent] = self.left[index] 149 | self.deleted += 1 150 | return True 151 | 152 | if self.has_right(index) and not self.has_left(index): 153 | if index == 0: 154 | let right = self.right[0] 155 | self.elements[0] = self.elements[right.to_int()] 156 | if self.has_left(right): 157 | self.left[0] = self.left[right.to_int()] 158 | else: 159 | self.left[0] = 0 160 | if self.has_right(right): 161 | self.right[0] = self.right[right.to_int()] 162 | else: 163 | self.right[0] = 0 164 | else: 165 | if self.left[parent] == index: 166 | self.left[parent] = self.right[index] 167 | else: 168 | self.right[parent] = self.right[index] 169 | self.deleted += 1 170 | return True 171 | 172 | return self._swap_with_next_smaller_leaf(index) 173 | 174 | @always_inline("nodebug") 175 | fn sorted_elements(self) -> UnsafeFixedVector[T]: 176 | let number_of_elements = self.__len__() 177 | var result = UnsafeFixedVector[T](number_of_elements) 178 | if number_of_elements == 0: 179 | return result 180 | var stack = DynamicVector[UInt16](self.max_depth.to_int()) 181 | var current: UInt16 = 0 182 | while len(result) < number_of_elements: 183 | if len(result) == 0 or cmp(result[len(result) - 1], self.elements[self.left[current.to_int()].to_int()]) < 0: 184 | while self.has_left(current): 185 | stack.push_back(current) 186 | current = self.left[current.to_int()] 187 | result.append(self.elements[current.to_int()]) 188 | if self.has_right(current): 189 | current = self.right[current.to_int()] 190 | else: 191 | current = stack.pop_back() 192 | 193 | return result 194 | 195 | fn clear(inout self): 196 | self.elements.clear() 197 | self.left.clear() 198 | self.right.clear() 199 | self.deleted = 0 200 | self.max_depth = 0 201 | self.balanced = False 202 | 203 | fn union(self, other: Self) -> Self: 204 | var result = Self() 205 | let combined: UnsafeFixedVector[T] 206 | if other.__len__() == 0: 207 | combined = self.sorted_elements() 208 | elif self.__len__() == 0: 209 | combined = other.sorted_elements() 210 | else: 211 | combined = self._combine[Self.Union](other) 212 | result._balance_with(combined) 213 | return result^ 214 | 215 | fn union_inplace(inout self, other: Self): 216 | if other.__len__() == 0: 217 | return 218 | if self.__len__() == 0: 219 | self._balance_with(other.sorted_elements()) 220 | return 221 | let combined = self._combine[Self.Union](other) 222 | self._balance_with(combined) 223 | 224 | fn intersection(self, other: Self) -> Self: 225 | var result = FibyTree[T, cmp, to_str]() 226 | if other.__len__() == 0: 227 | return result^ 228 | if self.__len__() == 0: 229 | return result^ 230 | let combined = self._combine[Self.Intersection](other) 231 | result._balance_with(combined) 232 | return result^ 233 | 234 | fn intersection_inplace(inout self, other: Self): 235 | if other.__len__() == 0: 236 | self.clear() 237 | return 238 | if self.__len__() == 0: 239 | self.clear() 240 | return 241 | let combined = self._combine[Self.Intersection](other) 242 | self._balance_with(combined) 243 | 244 | fn difference(self, other: Self) -> Self: 245 | var result = FibyTree[T, cmp, to_str]() 246 | let combined: UnsafeFixedVector[T] 247 | if other.__len__() == 0 or self.__len__() == 0: 248 | combined = self.sorted_elements() 249 | else: 250 | combined = self._combine[Self.Difference](other) 251 | result._balance_with(combined) 252 | return result^ 253 | 254 | fn difference_inplace(inout self, other: Self): 255 | if other.__len__() == 0 or self.__len__() == 0: 256 | return 257 | let combined = self._combine[Self.Difference](other) 258 | self._balance_with(combined) 259 | 260 | fn other_difference_inplace(inout self, other: Self): 261 | if other.__len__() == 0: 262 | self.clear() 263 | return 264 | if self.__len__() == 0: 265 | self._balance_with(other.sorted_elements()) 266 | return 267 | let combined = self._combine[Self.OtherDifference](other) 268 | self._balance_with(combined) 269 | 270 | fn symmetric_difference(self, other: Self) -> Self: 271 | var result = FibyTree[T, cmp, to_str]() 272 | let combined: UnsafeFixedVector[T] 273 | if other.__len__() == 0: 274 | combined = self.sorted_elements() 275 | elif self.__len__() == 0: 276 | combined = other.sorted_elements() 277 | else: 278 | combined = self._combine[Self.SymetricDifference](other) 279 | result._balance_with(combined) 280 | return result^ 281 | 282 | fn symmetric_difference_inplace(inout self, other: Self): 283 | if other.__len__() == 0: 284 | return 285 | if self.__len__() == 0: 286 | self._balance_with(other.sorted_elements()) 287 | return 288 | let combined = self._combine[Self.SymetricDifference](other) 289 | self._balance_with(combined) 290 | 291 | @always_inline("nodebug") 292 | fn _combine[type: Int](self, other: Self) -> UnsafeFixedVector[T]: 293 | let num1 = self.__len__() 294 | let num2 = other.__len__() 295 | # assert(num1 > 0) 296 | # assert(num2 > 0) 297 | var combined = UnsafeFixedVector[T](num1 + num2) 298 | var cur1: UInt16 = 0 299 | var cur2: UInt16 = 0 300 | var stack1 = DynamicVector[UInt16](self.max_depth.to_int()) 301 | var stack2 = DynamicVector[UInt16](other.max_depth.to_int()) 302 | var last_returned1 = UnsafeFixedVector[T](1) 303 | var last_returned2 = UnsafeFixedVector[T](1) 304 | var e1 = self._sorted_iter(cur1, stack1, last_returned1) 305 | last_returned1.append(e1) 306 | var e2 = other._sorted_iter(cur2, stack2, last_returned2) 307 | last_returned2.append(e2) 308 | var compute1 = False 309 | var compute2 = False 310 | var cursor1 = 1 311 | var cursor2 = 1 312 | var increase1 = False 313 | var increase2 = False 314 | while True: 315 | if compute1 and cursor1 < num1: 316 | e1 = self._sorted_iter(cur1, stack1, last_returned1) 317 | last_returned1.clear() 318 | last_returned1.append(e1) 319 | increase1 = True 320 | if compute2 and cursor2 < num2: 321 | e2 = other._sorted_iter(cur2, stack2, last_returned2) 322 | last_returned2.clear() 323 | last_returned2.append(e2) 324 | increase2 = True 325 | let diff = cmp(e1, e2) 326 | if diff < 0: 327 | if num1 == 1 and num2 == 1: 328 | @parameter 329 | if type == Self.Union or type == Self.Difference or type == Self.SymetricDifference: 330 | combined.append(e1) 331 | @parameter 332 | if type == Self.Union or type == Self.SymetricDifference or type == Self.OtherDifference: 333 | combined.append(e2) 334 | break 335 | if cursor1 < num1: 336 | @parameter 337 | if type == Self.Union or type == Self.Difference or type == Self.SymetricDifference: 338 | if len(combined) == 0 or cmp(combined[len(combined) - 1], e1) < 0: 339 | combined.append(e1) 340 | compute1 = cursor1 < num1 341 | compute2 = False 342 | else: 343 | @parameter 344 | if type == Self.Union or type == Self.SymetricDifference or type == Self.OtherDifference: 345 | if len(combined) == 0 or cmp(combined[len(combined) - 1], e2) < 0: 346 | combined.append(e2) 347 | compute1 = False 348 | compute2 = cursor2 < num2 349 | elif diff > 0: 350 | if num1 == 1 and num2 == 1: 351 | @parameter 352 | if type == Self.Union or type == Self.SymetricDifference or type == Self.OtherDifference: 353 | combined.append(e2) 354 | @parameter 355 | if type == Self.Union or type == Self.Difference or type == Self.SymetricDifference: 356 | combined.append(e1) 357 | break 358 | if cursor2 < num2: 359 | @parameter 360 | if type == Self.Union or type == Self.SymetricDifference or type == Self.OtherDifference: 361 | if len(combined) == 0 or cmp(combined[len(combined) - 1], e2) < 0: 362 | combined.append(e2) 363 | compute1 = False 364 | compute2 = cursor2 < num2 365 | else: 366 | @parameter 367 | if type == Self.Union or type == Self.Difference or type == Self.SymetricDifference: 368 | if len(combined) == 0 or cmp(combined[len(combined) - 1], e1) < 0: 369 | combined.append(e1) 370 | compute1 = cursor1 < num1 371 | compute2 = False 372 | else: 373 | @parameter 374 | if type == Self.Union or type == Self.Intersection: 375 | if len(combined) == 0 or cmp(combined[len(combined) - 1], e1) < 0: 376 | combined.append(e1) 377 | compute1 = cursor1 < num1 378 | compute2 = cursor2 < num2 379 | 380 | if increase1 and cursor1 < num1: 381 | cursor1 += 1 382 | increase1 = False 383 | if increase2 and cursor2 < num2: 384 | cursor2 += 1 385 | increase2 = False 386 | @parameter 387 | if type == Self.Intersection: 388 | if cursor1 >= num1 or cursor2 >= num2: 389 | break 390 | else: 391 | if cursor1 >= num1 and cursor2 >= num2: 392 | break 393 | 394 | return combined 395 | 396 | fn is_subset(self, other: Self) -> Bool: 397 | return self._check[Self.IsSubset](other) 398 | 399 | fn is_superset(self, other: Self) -> Bool: 400 | return self._check[Self.IsSuperset](other) 401 | 402 | fn is_disjoint(self, other: Self) -> Bool: 403 | return self._check[Self.IsDisjoint](other) 404 | 405 | @always_inline("nodebug") 406 | fn _check[type: Int](self, other: Self) -> Bool: 407 | let num1 = self.__len__() 408 | let num2 = other.__len__() 409 | @parameter 410 | if type == Self.IsSubset: 411 | if num1 == 0: 412 | return True 413 | if num1 > num2 or num2 == 0: 414 | return False 415 | @parameter 416 | if type == Self.IsSuperset: 417 | if num2 == 0: 418 | return True 419 | if num1 < num2 or num1 == 0: 420 | return False 421 | 422 | @parameter 423 | if type == Self.IsDisjoint: 424 | if num1 == 0 or num2 == 0: 425 | return True 426 | 427 | var cur1: UInt16 = 0 428 | var cur2: UInt16 = 0 429 | var stack1 = DynamicVector[UInt16](self.max_depth.to_int()) 430 | var stack2 = DynamicVector[UInt16](other.max_depth.to_int()) 431 | var last_returned1 = UnsafeFixedVector[T](1) 432 | var last_returned2 = UnsafeFixedVector[T](1) 433 | var e1 = self._sorted_iter(cur1, stack1, last_returned1) 434 | last_returned1.append(e1) 435 | var e2 = other._sorted_iter(cur2, stack2, last_returned2) 436 | last_returned2.append(e2) 437 | var compute1 = False 438 | var compute2 = False 439 | var cursor1 = 1 440 | var cursor2 = 1 441 | var increase1 = False 442 | var increase2 = False 443 | var num_eq = 0 444 | while True: 445 | if compute1 and cursor1 < num1: 446 | e1 = self._sorted_iter(cur1, stack1, last_returned1) 447 | last_returned1.clear() 448 | last_returned1.append(e1) 449 | increase1 = True 450 | if compute2 and cursor2 < num2: 451 | e2 = other._sorted_iter(cur2, stack2, last_returned2) 452 | last_returned2.clear() 453 | last_returned2.append(e2) 454 | increase2 = True 455 | let diff = cmp(e1, e2) 456 | if diff == 0: 457 | @parameter 458 | if type == Self.IsDisjoint: 459 | return False 460 | compute1 = cursor1 < num1 461 | compute2 = cursor2 < num2 462 | num_eq += 1 463 | else: 464 | if diff < 0: 465 | @parameter 466 | if type == Self.IsSubset: 467 | break 468 | compute1 = True 469 | compute2 = cursor1 >= num1 470 | else: 471 | @parameter 472 | if type == Self.IsSuperset: 473 | break 474 | compute1 = cursor2 >= num2 475 | compute2 = True 476 | 477 | if increase1 and cursor1 < num1: 478 | cursor1 += 1 479 | increase1 = False 480 | if increase2 and cursor2 < num2: 481 | cursor2 += 1 482 | increase2 = False 483 | 484 | if cursor1 >= num1 and cursor2 >= num2: 485 | break 486 | 487 | @parameter 488 | if type == Self.IsSuperset: 489 | return num_eq == num2 490 | @parameter 491 | if type == Self.IsSubset: 492 | return num_eq == num1 493 | @parameter 494 | if type == Self.IsDisjoint: 495 | return True 496 | return False 497 | 498 | @always_inline("nodebug") 499 | fn _sorted_iter(self, inout current: UInt16, inout stack: DynamicVector[UInt16], inout last_returned: UnsafeFixedVector[T]) -> T: 500 | # using UnsafeFixedVector[T](1) as poor mans Optional for last_returned 501 | if len(last_returned) == 0 or cmp(last_returned[0], self.elements[self.left[current.to_int()].to_int()]) < 0: 502 | while self.has_left(current): 503 | stack.push_back(current) 504 | current = self.left[current.to_int()] 505 | let result = self.elements[current.to_int()] 506 | if self.has_right(current): 507 | current = self.right[current.to_int()] 508 | else: 509 | current = stack.pop_back() 510 | return result 511 | 512 | @always_inline("nodebug") 513 | fn __len__(self) -> Int: 514 | return len(self.elements) - self.deleted 515 | 516 | @always_inline("nodebug") 517 | fn __contains__(self, element: T) -> Bool: 518 | return self._get_index(element).get[1, Int]() > -1 519 | 520 | fn _get_index(self, element: T) -> (Int, Int): 521 | if self.__len__() == 0: 522 | return -1, -1 523 | if self.balanced: 524 | return self._get_index_balanced(element) 525 | var parent = 0 526 | var index = 0 527 | while True: 528 | let diff = cmp(self.elements[index], element) 529 | if diff == 0: 530 | return parent, index 531 | if diff > 0: 532 | let left = self.left[index].to_int() 533 | if left == index: 534 | return index, -1 535 | else: 536 | parent = index 537 | index = left 538 | else: 539 | let right = self.right[index].to_int() 540 | if right == index: 541 | return index, -1 542 | else: 543 | parent = index 544 | index = right 545 | 546 | fn _get_index_balanced(self, element: T) -> (Int, Int): 547 | var parent = 0 548 | var index = 0 549 | let len = self.__len__() 550 | while index < len: 551 | let diff = cmp(element, self.elements[index]) 552 | if diff == 0: 553 | return parent, index 554 | parent = index 555 | index = (index + 1) * 2 + (diff >> 63) 556 | return parent, -1 557 | 558 | fn min_index(self) -> Int: 559 | if self.__len__() < 2: 560 | return self.__len__() - 1 561 | if self.balanced: 562 | return (1 << (self.max_depth.to_int() - 1)) - 1 563 | var cand = self.left[0] 564 | while self.has_left(cand): 565 | cand = self.left[cand.to_int()] 566 | return cand.to_int() 567 | 568 | fn max_index(self) -> Int: 569 | let size = self.__len__() 570 | if size < 2: 571 | return size - 1 572 | if self.balanced: 573 | if size == (1 << self.max_depth.to_int()) - 1: 574 | return size - 1 575 | return (1 << (self.max_depth.to_int() - 1)) - 2 576 | var cand = self.right[0] 577 | while self.has_right(cand): 578 | cand = self.right[cand.to_int()] 579 | return cand.to_int() 580 | 581 | 582 | fn _swap_with_next_smaller_leaf(inout self, index: UInt16) -> Bool: 583 | var parent = index 584 | var candidate = self.left[index.to_int()] 585 | if candidate == index: 586 | return False 587 | while True: 588 | if self._is_leaf(candidate): 589 | self.elements[index.to_int()] = self.elements[candidate.to_int()] 590 | self._delete_leaf(candidate.to_int(), parent.to_int()) 591 | return True 592 | let right = self.right[candidate.to_int()] 593 | if right == candidate: 594 | self.elements[index.to_int()] = self.elements[candidate.to_int()] 595 | self.right[parent.to_int()] = self.left[candidate.to_int()] 596 | self.deleted += 1 597 | return True 598 | else: 599 | parent = candidate 600 | candidate = right 601 | 602 | @always_inline("nodebug") 603 | fn _is_leaf(self, index: UInt16) -> Bool: 604 | return (self.left[index.to_int()] == index).__bool__() and (self.right[index.to_int()] == index).__bool__() 605 | 606 | @always_inline("nodebug") 607 | fn _delete_leaf(inout self, index: Int, parent: Int): 608 | self.deleted += 1 609 | if self.left[parent] == index: 610 | self.left[parent] = parent 611 | else: 612 | self.right[parent] = parent 613 | 614 | fn balance(inout self): 615 | if self.balanced: 616 | return 617 | let sorted_elements = self.sorted_elements() 618 | self._balance_with(sorted_elements) 619 | 620 | @always_inline("nodebug") 621 | fn _balance_with(inout self, sorted_elements: UnsafeFixedVector[T]): 622 | let new_size = len(sorted_elements) 623 | self.elements.resize(new_size) 624 | self.left.resize(new_size) 625 | self.right.resize(new_size) 626 | 627 | var i: Int = 0 628 | self._eytzinger(i, 1, sorted_elements) 629 | for index in range(new_size): 630 | let l = (index + 1) * 2 - 1 631 | let r = (index + 1) * 2 632 | if l < self.__len__(): 633 | self.left[index] = l 634 | else: 635 | self.left[index] = index 636 | if r < self.__len__(): 637 | self.right[index] = r 638 | else: 639 | self.right[index] = index 640 | 641 | self.deleted = 0 642 | 643 | self.balanced = True 644 | self.max_depth = self._optimal_depth() 645 | 646 | fn _eytzinger(inout self, inout i: Int, k: Int, v: UnsafeFixedVector[T]): 647 | if k <= len(v): 648 | self._eytzinger(i, k * 2, v) 649 | self.elements[k - 1] = v[i] 650 | i += 1 651 | self._eytzinger(i, k * 2 + 1, v) 652 | 653 | fn print_tree(self, root: UInt16 = 0): 654 | if self.__len__() == 0: 655 | print("・") 656 | return 657 | self._print("", 0) 658 | 659 | fn _print(self, indentation: String, index: UInt16): 660 | if len(indentation) > 0: 661 | print(indentation, "-", to_str(self.elements[index.to_int()])) 662 | else: 663 | print("-", to_str(self.elements[index.to_int()])) 664 | 665 | if self.has_left(index): 666 | self._print(indentation + " ", self.left[index.to_int()]) 667 | elif self.has_right(index): 668 | print(indentation + " ", "- ・") 669 | if self.has_right(index): 670 | self._print(indentation + " ", self.right[index.to_int()]) --------------------------------------------------------------------------------