1classdef BalancedBST < handle
2 %BALANCEDBST AVL-balanced binary search tree mapping keys to data. See
3 % https://en.wikipedia.org/wiki/AVL_tree for further infos.
4 %
5 % Keys are typically char vectors (strings); data is typically a
6 % scalar, but both can be any MATLAB type whose comparison is
7 % well-defined (numeric or char/string for keys).
8 %
9 % Performance notes
10 % -----------------
11 % Cell arrays in MATLAB are contiguous pointer arrays with O(1)
12 % random access — NOT linked lists. The tree traversal is
13 % O(log n) key-comparisons. The pre-allocated parallel-array
14 % storage avoids per-node heap allocations; doubling on overflow
15 % gives amortised O(1) growth.
16 %
17 % Usage
18 % -----
19 % tree = mr.aux.BalancedBST();
20 % tree.insert('alpha', 1);
21 % tree.insert('beta', 2);
22 %
23 % [val, index] = tree.lookup('alpha', 0); % val=1, index=non-zero
24 % [val, index] = tree.lookup('gamma', 0); % val=0, index=0
25 %
27 % --- storage: parallel arrays (pre-allocated, grown as needed) -------
28 properties (Access = private)
29 keys % cell(1,cap) – stored keys
30 vals % cell(1,cap) – stored data values
31 L % int32(1,cap) – left-child index (0 = none)
32 R % int32(1,cap) – right-child index (0 = none)
33 H % int32(1,cap) – subtree height
34 rootIdx % int32 – root node index (0 = empty tree)
35 cnt % int32 – number of allocated slots (high-water mark)
36 cap % int32 – allocated capacity
37 keyIsChar % bool – whether the key is a char string or an arbitrary vector
38 pathBuf % int32
39 dirBuf % int32
40 freeHead % int32 – head of free-slot linked list (0 = none)
41 liveCnt % int32 – number of live (non-deleted) nodes
42 end
44 % =====================================================================
45 % PUBLIC INTERFACE
46 % =====================================================================
47 methods (Access = public)
49 function obj = BalancedBST(initialCapacity)
50 %BALANCEDBST Construct an empty tree.
51 % tree = mr.aux.BalancedBST()
52 % tree = mr.aux.BalancedBST(initialCapacity)
53 if nargin < 1, initialCapacity = 64; end
54 obj.cap = int32(initialCapacity);
55 obj.keys = cell(1, obj.cap);
56 obj.vals = cell(1, obj.cap);
57 obj.L = zeros(1, obj.cap, 'int32');
58 obj.R = zeros(1, obj.cap, 'int32');
59 obj.H = zeros(1, obj.cap, 'int32');
60 obj.rootIdx = int32(0);
61 obj.cnt = int32(0);
62 obj.keyIsChar = false;
63 obj.pathBuf = int32(0);
64 obj.dirBuf = int32(0);
65 obj.freeHead = int32(0);
66 obj.liveCnt = int32(0);
67 end
69 function [val, index] = lookup(obj, key, default)
70 %LOOKUP Search for KEY in the tree.
71 % [val, index] = tree.lookup(key, default)
72 %
73 % If KEY is present, VAL is the associated data and
74 % INDEX is non-zero. Otherwise VAL = DEFAULT and INDEX is
75 % zero.
76 %
78 val = default;
79 index = 0;
81 % --- cache property arrays as locals (COW = free) ------------
82 % Every obj.prop(idx) in a loop pays ~300 ns dispatch
83 % overhead; local variable access is ~1 ns. For a
84 % read-only traversal COW means the snapshot is free.
85 %kk = obj.keys;
86 %LL = obj.L;
87 %RR = obj.R;
89 idx = obj.rootIdx;
90 while idx ~= int32(0)
91 c = mr.aux.BalancedBST.compareKeys(key, obj.keys{idx});
92 if c == 0
93 val = obj.vals{idx}; % single property access
94 index = idx;
95 break;
96 end
97 if c < 0
98 idx = obj.L(idx);
99 else
100 idx = obj.R(idx);
101 end
102 end
103 end
105 function insert(obj, key, val)
106 %INSERT Insert or update a key-value pair.
107 % tree.insert(key, val)
108 %
109 % If KEY already exists its data is overwritten.
111 root = obj.rootIdx;
112 if root == int32(0) % the tree was empty up to now
113 obj.rootIdx = obj.allocNode(key, val);
114 obj.keyIsChar = ischar(key);
115 return;
116 end
118 % --- search phase --------------------------------------------
119 depth = int32(0);
120 idx = root;
122 while idx ~= int32(0)
123 c = mr.aux.BalancedBST.compareKeys(key, obj.keys{idx});
124 depth = depth + 1;
125 obj.pathBuf(depth) = idx;
126 if c == 0
127 obj.vals{idx} = val; % overwrite existing value
128 return;
129 elseif c < 0
130 obj.dirBuf(depth) = int32(-1);
131 idx = obj.L(idx);
132 else
133 obj.dirBuf(depth) = int32(1);
134 idx = obj.R(idx);
135 end
136 end
138 % --- allocate new node and link to parent --------------------
139 newIdx = obj.allocNode(key, val);
140 pIdx = obj.pathBuf(depth);
141 if obj.dirBuf(depth) < 0
142 obj.L(pIdx) = newIdx;
143 else
144 obj.R(pIdx) = newIdx;
145 end
147 % --- rebalance bottom-up -------------------------------------
148 obj.rebalanceUp(depth);
149 end
151 function update(obj, index, val)
152 %UPDATE Update the value at a known node index in O(1).
153 % tree.update(index, newVal)
154 %
155 % INDEX is typically obtained from a prior tree.lookup().
156 % The key is unchanged; only the associated data is
157 % overwritten. No rebalancing is needed.
158 obj.vals{index} = val;
159 end
161 function removed = remove(obj, key)
162 %REMOVE Remove the entry with the given key.
163 % removed = tree.remove(key)
164 %
165 % Returns true if KEY was found and removed, false if
166 % KEY was not present in the tree.
168 if obj.rootIdx == int32(0)
169 removed = false;
170 return;
171 end
173 % --- search for the node -------------------------------------
174 depth = int32(0);
175 idx = obj.rootIdx;
176 while idx ~= int32(0)
177 c = mr.aux.BalancedBST.compareKeys(key, obj.keys{idx});
178 depth = depth + 1;
179 obj.pathBuf(depth) = idx;
180 if c == 0
181 break;
182 elseif c < 0
183 obj.dirBuf(depth) = int32(-1);
184 idx = obj.L(idx);
185 else
186 obj.dirBuf(depth) = int32(1);
187 idx = obj.R(idx);
188 end
189 end
191 if idx == int32(0)
192 removed = false;
193 return;
194 end
196 % idx == obj.pathBuf(depth) is the node to delete
197 li = obj.L(idx);
198 ri = obj.R(idx);
200 if li ~= 0 && ri ~= 0
201 % --- TWO CHILDREN: replace with in-order successor --------
202 % Go right once, then left as far as possible.
203 obj.dirBuf(depth) = int32(1); % going right from idx
204 succ = ri;
205 depth = depth + 1;
206 obj.pathBuf(depth) = succ;
207 while obj.L(succ) ~= int32(0)
208 obj.dirBuf(depth) = int32(-1);
209 succ = obj.L(succ);
210 depth = depth + 1;
211 obj.pathBuf(depth) = succ;
212 end
213 % Copy successor's key/value to the target node
214 obj.keys{idx} = obj.keys{succ};
215 obj.vals{idx} = obj.vals{succ};
216 % Successor has at most a right child
217 replacement = obj.R(succ);
218 obj.freeNode(succ);
219 else
220 % --- ZERO or ONE CHILD -----------------------------------
221 if li ~= 0
222 replacement = li;
223 else
224 replacement = ri; % may be 0 (leaf)
225 end
226 obj.freeNode(idx);
227 end
229 % --- link replacement to parent of deleted node ---------------
230 if depth > 1
231 p = obj.pathBuf(depth - 1);
232 if obj.dirBuf(depth - 1) < 0
233 obj.L(p) = replacement;
234 else
235 obj.R(p) = replacement;
236 end
237 else
238 obj.rootIdx = replacement;
239 end
241 % --- rebalance from parent of deleted node upward -------------
242 obj.rebalanceUp(depth - 1);
244 removed = true;
245 end
247 function n = length(obj)
248 %LENGTH Number of live key-value pairs in the tree.
249 n = double(obj.liveCnt);
250 end
252 % function varargout = subsref(obj, S)
253 % %SUBSREF Overloaded subscript reference.
254 % % val = tree('key') — equivalent to tree.lookup('key', [])
255 % %
256 % % Dot-reference (tree.method, tree.prop) and curly-brace
257 % % indexing are forwarded to the built-in handler so that
258 % % normal method calls keep working.
259 % if S(1).type(1) == '('
260 % key = S(1).subs{1};
261 % [val, ~] = obj.lookup(key, []);
262 % if numel(S) > 1
263 % % chained indexing, e.g. tree('key').field
264 % [varargout{1:nargout}] = subsref(val, S(2:end));
265 % else
266 % varargout{1} = val;
267 % end
268 % else
269 % % '.' or '{}' — delegate to built-in
270 % [varargout{1:nargout}] = builtin('subsref', obj, S);
271 % end
272 % end
274 function obj = subsasgn(obj, S, val)
275 %SUBSASGN Overloaded subscript assignment.
276 % tree('key') = val — equivalent to tree.insert('key', val)
277 %
278 % Dot-assignment and curly-brace assignment are forwarded
279 % to the built-in handler.
280 if S(1).type(1) == '(' && numel(S) == 1
281 key = S(1).subs{1};
282 obj.insert(key, val);
283 else
284 % '.' or '{}' or chained — delegate to built-in
285 obj = builtin('subsasgn', obj, S, val);
286 end
287 end
289 end % public methods
291 % =====================================================================
292 % PRIVATE HELPERS
293 % =====================================================================
294 methods (Access = private)
296 % ----- node allocation -------------------------------------------
297 function idx = allocNode(obj, key, val)
298 if obj.freeHead ~= int32(0)
299 idx = obj.freeHead;
300 obj.freeHead = obj.L(idx); % L was reused as next-free
301 else
302 obj.cnt = obj.cnt + 1;
303 if obj.cnt > obj.cap
304 obj.grow();
305 end
306 idx = obj.cnt;
307 end
308 obj.keys{idx} = key;
309 obj.vals{idx} = val;
310 obj.L(idx) = int32(0);
311 obj.R(idx) = int32(0);
312 obj.H(idx) = int32(1);
313 obj.liveCnt = obj.liveCnt + 1;
314 end
316 function freeNode(obj, idx)
317 %FREENODE Return a slot to the free list.
318 obj.keys{idx} = [];
319 obj.vals{idx} = [];
320 obj.R(idx) = int32(0);
321 obj.H(idx) = int32(0);
322 obj.L(idx) = obj.freeHead; % reuse L as next-free pointer
323 obj.freeHead = idx;
324 obj.liveCnt = obj.liveCnt - 1;
325 end
327 function grow(obj)
328 added = obj.cap; % double the capacity
329 obj.keys = [obj.keys, cell(1, added)];
330 obj.vals = [obj.vals, cell(1, added)];
331 obj.L = [obj.L, zeros(1, added, 'int32')];
332 obj.R = [obj.R, zeros(1, added, 'int32')];
333 obj.H = [obj.H, zeros(1, added, 'int32')];
334 obj.cap = obj.cap + int32(added);
335 end
337 % ----- rebalance from pathBuf(depth) up to root ------------------
338 function rebalanceUp(obj, depth)
339 for i = depth:-1:1
340 nd = obj.pathBuf(i);
342 % -- refresh height (inlined) --
343 li = obj.L(nd); ri = obj.R(nd);
344 lh = int32(0); rh = int32(0);
345 if li ~= 0, lh = obj.H(li); end
346 if ri ~= 0, rh = obj.H(ri); end
347 obj.H(nd) = int32(1) + max(lh, rh);
349 bf = rh - lh;
350 nnd = nd;
352 if bf < -1
353 % left-heavy
354 child = obj.L(nd);
355 cli = obj.L(child); cri = obj.R(child);
356 clh = int32(0); crh = int32(0);
357 if cli ~= 0, clh = obj.H(cli); end
358 if cri ~= 0, crh = obj.H(cri); end
359 if (crh - clh) > 0 % Left-Right case
360 gc = obj.R(child);
361 obj.R(child) = obj.L(gc);
362 obj.L(gc) = child;
363 tl = obj.L(child); tr = obj.R(child);
364 tlh = int32(0); trh = int32(0);
365 if tl ~= 0, tlh = obj.H(tl); end
366 if tr ~= 0, trh = obj.H(tr); end
367 obj.H(child) = int32(1) + max(tlh, trh);
368 tl = obj.L(gc); tr = obj.R(gc);
369 tlh = int32(0); trh = int32(0);
370 if tl ~= 0, tlh = obj.H(tl); end
371 if tr ~= 0, trh = obj.H(tr); end
372 obj.H(gc) = int32(1) + max(tlh, trh);
373 obj.L(nd) = gc;
374 end
375 x = obj.L(nd);
376 obj.L(nd) = obj.R(x);
377 obj.R(x) = nd;
378 tl = obj.L(nd); tr = obj.R(nd);
379 tlh = int32(0); trh = int32(0);
380 if tl ~= 0, tlh = obj.H(tl); end
381 if tr ~= 0, trh = obj.H(tr); end
382 obj.H(nd) = int32(1) + max(tlh, trh);
383 tl = obj.L(x); tr = obj.R(x);
384 tlh = int32(0); trh = int32(0);
385 if tl ~= 0, tlh = obj.H(tl); end
386 if tr ~= 0, trh = obj.H(tr); end
387 obj.H(x) = int32(1) + max(tlh, trh);
388 nnd = x;
390 elseif bf > 1
391 % right-heavy
392 child = obj.R(nd);
393 cli = obj.L(child); cri = obj.R(child);
394 clh = int32(0); crh = int32(0);
395 if cli ~= 0, clh = obj.H(cli); end
396 if cri ~= 0, crh = obj.H(cri); end
397 if (crh - clh) < 0 % Right-Left case
398 gc = obj.L(child);
399 obj.L(child) = obj.R(gc);
400 obj.R(gc) = child;
401 tl = obj.L(child); tr = obj.R(child);
402 tlh = int32(0); trh = int32(0);
403 if tl ~= 0, tlh = obj.H(tl); end
404 if tr ~= 0, trh = obj.H(tr); end
405 obj.H(child) = int32(1) + max(tlh, trh);
406 tl = obj.L(gc); tr = obj.R(gc);
407 tlh = int32(0); trh = int32(0);
408 if tl ~= 0, tlh = obj.H(tl); end
409 if tr ~= 0, trh = obj.H(tr); end
410 obj.H(gc) = int32(1) + max(tlh, trh);
411 obj.R(nd) = gc;
412 end
413 y = obj.R(nd);
414 obj.R(nd) = obj.L(y);
415 obj.L(y) = nd;
416 tl = obj.L(nd); tr = obj.R(nd);
417 tlh = int32(0); trh = int32(0);
418 if tl ~= 0, tlh = obj.H(tl); end
419 if tr ~= 0, trh = obj.H(tr); end
420 obj.H(nd) = int32(1) + max(tlh, trh);
421 tl = obj.L(y); tr = obj.R(y);
422 tlh = int32(0); trh = int32(0);
423 if tl ~= 0, tlh = obj.H(tl); end
424 if tr ~= 0, trh = obj.H(tr); end
425 obj.H(y) = int32(1) + max(tlh, trh);
426 nnd = y;
427 end
429 % -- re-link to parent ------------------------------------
430 if nnd ~= nd
431 if i > 1
432 p = obj.pathBuf(i-1);
433 if obj.dirBuf(i-1) < 0
434 obj.L(p) = nnd;
435 else
436 obj.R(p) = nnd;
437 end
438 else
439 obj.rootIdx = nnd;
440 end
441 end
442 end
443 end
445 end % private methods
447 % =====================================================================
448 % STATIC (key comparison)
449 % =====================================================================
450 methods (Static, Access = private)
452 function c = compareKeys(a, b)
453 %COMPAREKEYS Lexicographic comparison returning -1, 0, or +1.
454 % Handles numeric keys (scalar <, >, ==) and char/string
455 % keys (character-by-character comparison).
456 % [~,I]=sort({a,b});
457 % c = diff(I)*~strcmp(a,b);
458 %la = numel(a); lb = numel(b);
459 %ml = min(la, lb);
460 for k = 1:min(numel(a),numel(b))
461 % c=sign(int32(a(k))-int32(b(k)));
462 % if c~=0
463 % return;
464 % end
465 if a(k) > b(k)
466 c = 1;
467 return;
468 elseif a(k) < b(k)
469 c = -1;
470 return;
471 end
472 end
473 c = sign(numel(a)-numel(b));
474 % a=string(a);
475 % b=string(b);
476 % if a>b
477 % c = 1;
478 % elseif a < b
479 % c = -1;
480 % else
481 % c = 0;
482 % end
483 end
485 end % static methods
486end