• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1/*
2 * Copyright (c) 2021-2024 Huawei Device Co., Ltd.
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16package std.containers;
17
18/**
19 * Node for a tree
20 */
21export class Node<K extends Comparable<Object>, V> {
22    left: Node<K, V> | null;
23    key: K;
24    value: V;
25    height: int;
26    right: Node<K, V> | null;
27    constructor(left: Node<K, V> | null, key: K, value: V, height: int, right: Node<K, V> | null) {
28        this.key = key;
29        this.value = value;
30        this.left = left;
31        this.right = right;
32        this.height = height;
33    }
34}
35
36/**
37 * Constructs new empty Node
38 *
39 * @returns empty node
40 */
41export function newEmpty<K extends Comparable<Object>, V>(): Node<K, V> | null {
42    return null;
43}
44
45/**
46 * Checks if Node is empty
47 *
48 * @param n
49 *
50 * @returns true if Node is empty, false otherwise
51 */
52export function isEmpty<K extends Comparable<Object>, V>(n: Node<K, V> | null): boolean {
53    return n == null;
54}
55
56/**
57 * Checks if Node is a leaf
58 *
59 * @param n
60 *
61 * @returns true if Node is a leaf, false otherwise
62 */
63export function isLeaf<K extends Comparable<Object>, V>(n: Node<K, V> | null): boolean {
64    return n != null && n!.left == null && n!.right == null;
65}
66
67/**
68 * Checks if Node is a leaf
69 *
70 * @param n
71 *
72 * @returns true if Node is a leaf, false otherwise
73 */
74export function isNode<K extends Comparable<Object>, V>(n: Node<K, V> | null): boolean {
75    return n != null && (n!.left != null || n!.right != null);
76}
77
78export function height<K extends Comparable<Object>, V>(node: Node<K, V> | null): int {
79    if (isEmpty<K, V>(node)) {
80        return 0;
81    } else if (isLeaf(node)) {
82        return 1;
83    } else {
84        return node!.height;
85    }
86}
87
88export function newNode<K extends Comparable<Object>, V>(left: Node<K, V> | null, k: K, v: V, height: int, right: Node<K, V> | null): Node<K, V> {
89    return new Node<K, V>(left, k, v, height, right);
90}
91
92export function newLeaf<K extends Comparable<Object>, V>(k: K, v: V): Node<K, V> {
93    return new Node<K, V>(newEmpty<K, V>(), k, v, 1, newEmpty<K, V>());
94}
95
96function legalLeftKey<K extends Comparable<Object>, V>(node: Node<K, V> | null, key: K): void {
97    if (isEmpty(node)) {
98        return;
99    }
100
101    let leftKey: K = node!.key;
102    assert leftKey.compareTo(key) < 0;
103}
104
105function legalRightKey<K extends Comparable<Object>, V>(node: Node<K, V> | null, key: K): void {
106    if (isEmpty(node)) {
107        return;
108    }
109
110    let rightKey: K = node!.key ;
111    assert rightKey.compareTo(key) > 0;
112}
113
114function inv<K extends Comparable<Object>, V>(node: Node<K, V> | null): void {
115    if (isNode(node)) {
116        let hl = height(node!.left);
117        let hr = height(node!.right);
118
119        inv(node!.left);
120        inv(node!.right);
121
122        legalLeftKey(node!.left, node!.key);
123        legalRightKey(node!.right, node!.key);
124
125        assert node!.height == max(hl, hr) + 1;
126        assert abs(hl - hr) <= 2;
127    }
128}
129
130
131function updateHeight<K extends Comparable<Object>, V>(node: Node<K, V> | null): void {
132    if (isNode(node)) {
133        let newHeight = max(height(node!.left), height(node!.right)) + 1;
134        if (newHeight != node!.height) {
135            node!.height = newHeight;
136        }
137    } else if (isLeaf(node)) {
138        node!.height = 1;
139    } else {
140        assert false;
141    }
142}
143
144function balanceTree<K extends Comparable<Object>, V>(node: Node<K, V> | null): Node<K, V> | null {
145    if (isEmpty(node) || isLeaf(node)) {
146        return node;
147    } else if (isNode(node)) {
148        let rootNode = node!;
149        let hl = height(rootNode.left);
150        let hr = height(rootNode.right);
151
152        if (hl > hr + 2) {
153            assert !(isEmpty(rootNode.left)) && !(isLeaf(rootNode.left));
154            let leftNode = rootNode.left!;
155            if (height(leftNode.left) >= height(leftNode.right)) {
156                rootNode.left = leftNode.right;
157                leftNode.right = rootNode;
158
159                updateHeight(rootNode);
160                updateHeight(leftNode);
161
162                return leftNode;
163            } else {
164                assert !(isEmpty(leftNode.right)) && !(isLeaf(leftNode.right));
165                let lrNode = leftNode.right!;
166
167                leftNode.right = lrNode.left;
168                rootNode.left = lrNode.right;
169                lrNode.right = rootNode;
170                lrNode.left = leftNode;
171
172                updateHeight(leftNode);
173                updateHeight(rootNode);
174                updateHeight(lrNode);
175
176                return lrNode;
177            }
178        } else if (hr > hl + 2) {
179            assert !(isEmpty(rootNode.right)) && !(isLeaf(rootNode.right));
180            let rightNode = rootNode.right!;
181            if (height(rightNode.right) >= height(rightNode.left)) {
182                rootNode.right = rightNode.left;
183                rightNode.left = rootNode;
184
185                updateHeight(rootNode);
186                updateHeight(rightNode);
187
188                return rightNode;
189            } else {
190                assert !(isEmpty(rightNode.left)) && !(isLeaf(rightNode.left));
191                let rlNode = rightNode.left!;
192
193                rightNode.left = rlNode.right;
194                rootNode.right = rlNode.left;
195                rlNode.left = rootNode;
196                rlNode.right = rightNode;
197
198                updateHeight(rightNode);
199                updateHeight(rootNode);
200                updateHeight(rlNode);
201
202                return rlNode;
203            }
204        } else {
205            updateHeight(node);
206            return node;
207        }
208    } else {
209        assert false;
210	return null;
211    }
212}
213
214function setLeft<K extends Comparable<Object>, V>(node: Node<K, V>, tree: Node<K, V> | null): void {
215    let tree1 = balanceTree(tree);
216    assert isNode(node);
217
218    let n = node;
219    if (n.left != tree1) {
220        n.left = tree1;
221    }
222
223    updateHeight(n);
224}
225
226function setRight<K extends Comparable<Object>, V>(node: Node<K, V>, tree: Node<K, V> | null): void {
227    let tree1 = balanceTree(tree);
228    assert isNode(node);
229
230    let n = node;
231    if (n.right != tree1) {
232        n.right = tree1;
233    }
234
235    updateHeight(n);
236}
237class addResult<K extends Comparable<Object>, V> {
238    t: Node<K, V>;
239    res: boolean;
240    constructor(t: Node<K, V>, res: boolean) {
241        this.t = t;
242        this.res = res;
243    }
244}
245
246function add<K extends Comparable<Object>, V>(node: Node<K, V> | null, k: K, v: V, replace: boolean): addResult<K, V> {
247    if (isEmpty(node)) {
248        return new addResult<K, V>(newLeaf(k, v), true);
249    }
250
251    if (isLeaf(node)) {
252        let leaf = node as Node<K, V>;
253        let c = leaf.key.compareTo(k);
254        if (c == 0) {
255            if (replace) {
256                leaf.value = v;
257                return new addResult<K, V>(leaf, false);
258            } else {
259                return new addResult<K, V>(leaf, false);
260            }
261        } else {
262            if (c < 0) {
263                return new addResult<K, V>(newNode(node, k, v, 2, newEmpty<K, V>()), true);
264            } else {
265                return new addResult<K, V>(newNode(newEmpty<K, V>(), k, v, 2, node), true);
266            }
267        }
268    }
269
270    if (isNode(node)) {
271        let nnode = node as Node<K, V>;
272        let c = k.compareTo(nnode.key);
273        if (c == 0) {
274            if (replace) {
275                nnode.value = v;
276            }
277            return new addResult<K, V>(nnode, false);
278        } else {
279            if (c < 0) {
280                let r: addResult<K, V> = add(nnode.left, k, v, replace);
281                setLeft(nnode, r.t);
282                return new addResult<K, V>(nnode, r.res);
283            } else {
284                let r: addResult<K, V> = add(nnode.right, k, v, replace);
285                setRight(nnode, r.t);
286                return new addResult<K, V>(nnode, r.res);
287            }
288        }
289    }
290
291    throw new Error("Invalid node");
292}
293
294
295export function addToTree<K extends Comparable<Object>, V>(node: Node<K, V> | null, k: K, v: V): Node<K, V> {
296    let r: addResult<K, V> = add(node, k, v, true);
297    if (!r.res) {
298        return balanceTree(r.t) as Node<K, V>
299    } else {
300        return r.t
301    }
302}
303
304function minElt<K extends Comparable<Object>, V>(node: Node<K, V> | null): Node<K, V> | null {
305    if (isEmpty(node) || isLeaf(node)) {
306        return node;
307    }
308    if (isEmpty(node!.left)) {
309        return node;
310    } else {
311        return minElt(node!.left);
312    }
313}
314
315function removeMinElt<K extends Comparable<Object>, V>(node: Node<K, V>): Node<K, V> | null {
316    assert !(isEmpty(node));
317    if (isLeaf(node)) {
318        return null;
319    }
320    if (isEmpty(node.left)) {
321        return node.right;
322    } else if (isLeaf(node.left) && isEmpty(node.right)) {
323        return newLeaf(node.key, node.value);
324    } else if (isLeaf(node.left)) {
325        setLeft(node, null);
326        return node;
327    } else {
328        setLeft(node, removeMinElt(node.left as Node<K, V>));
329        return node;
330    }
331}
332
333export function mergeTree<K extends Comparable<Object>, V>(t1: Node<K, V> | null, t2: Node<K, V> | null): Node<K, V> | null {
334    if (isEmpty(t1)) {
335        return t2;
336    }
337    if (!isEmpty(t1) && isEmpty(t2)) {
338        return t1;
339    }
340    let nt2 = t2 as Node<K, V>;
341
342    let tree: Node<K, V> = minElt(nt2) as Node<K, V>;
343
344    assert !isEmpty(tree);
345
346    if (isLeaf(tree)) {
347        let leaf = tree;
348        let t3 = balanceTree(removeMinElt(nt2));
349        return newNode(t1, leaf.key, leaf.value, max(height(t1), height(t3)) + 1, t3);
350    } else {
351        let node = tree;
352        setRight(node, removeMinElt(nt2));
353        setLeft(node, t1);
354        return node;
355    }
356}
357
358final class removeResult<K extends Comparable<Object>, V> {
359    t: Node<K, V> | null;
360    res: boolean;
361    constructor(t: Node<K, V> | null, res: boolean) {
362        this.t = t;
363        this.res = res;
364    }
365}
366
367function remove<K extends Comparable<Object>, V>(node: Node<K, V> | null, k: K): removeResult<K, V> {
368    if (isEmpty(node)) {
369        return new removeResult<K, V>(node, false);
370    } else if (isLeaf(node)) {
371        let leaf = node!;
372        if (leaf.key.compareTo(k) == 0) {
373            return new removeResult<K, V>(newEmpty<K, V>(), true);
374        } else {
375            return new removeResult<K, V>(node, false);
376        }
377    } else {
378        let nnode = node as Node<K, V>;
379        let c = k.compareTo(nnode.key);
380        if (c == 0) {
381            return new removeResult<K, V>(mergeTree(nnode.left, nnode.right), true);
382        } else if (c < 0) {
383            let r: removeResult<K, V> = remove(nnode.left, k);
384            setLeft(nnode, r.t)
385            return new removeResult<K, V>(nnode, r.res);
386        } else {
387            let r: removeResult<K, V> = remove(nnode.right, k);
388            setRight(nnode, r.t)
389            return new removeResult<K, V>(nnode, r.res);
390        }
391    }
392}
393
394export function removeFromTree<K extends Comparable<Object>, V>(node: Node<K, V> | null, k: K): Node<K, V> | null {
395    let r: removeResult<K, V> = remove(node, k);
396    return balanceTree(r.t);
397}
398
399/**
400 * Finds the specified value in the Node tree
401 *
402 * @param node tree root
403 *
404 * @param k value to find
405 *
406 * @returns value
407 * TODO: replace null
408 */
409export function lookupTree<K extends Comparable<Object>, V>(node: Node<K, V> | null, k: K): V | null {
410    if (isEmpty(node)) {
411        return null;
412    } else if (isLeaf(node)) {
413        let leaf = node!;
414        if (k.compareTo(leaf.key) == 0) {
415            return leaf.value;
416        } else {
417            return null;
418        }
419    } else {
420        let c = k.compareTo(node!.key);
421        if (c == 0) {
422            return node!.value;
423        } else if (c < 0) {
424            return lookupTree(node!.left, k);
425        } else {
426            return lookupTree(node!.right, k);
427        }
428    }
429}
430
431export function count<K extends Comparable<Object>, V>(node: Node<K, V> | null): int {
432    if (isEmpty(node)) {
433        return 0;
434    } else if (isLeaf(node)) {
435        return 1;
436    } else {
437        return 1 + count(node!.left) + count(node!.right);
438    }
439}
440
441function treeFolder<K extends Comparable<Object>, V>(k: K, v: V, acc: Entry<K, V>[]): Entry<K, V>[] {
442    let entry = new Entry<K, V>(k, v);
443    let acc1 = new Entry<K, V>[acc.length + 1];
444
445    for (let i: int = 0; i < acc.length; i++) {
446        acc1[i] = acc[i];
447    }
448
449    acc1[acc.length] = entry;
450
451    return acc1;
452}
453
454function foldTree<K extends Comparable<Object>, V>(node: Node<K, V> | null, acc0: Entry<K, V>[], fn: (k: K, v: V, acc: Entry<K, V>[]) => Entry<K, V>[]): Entry<K, V>[] {
455    if (isEmpty(node)) {
456        return acc0;
457    } else {
458        let acc1 = foldTree(node!.left, acc0, fn);
459        let acc2 = fn(node!.key, node!.value, acc1);
460        return foldTree(node!.right, acc2, fn);
461    }
462}
463
464export function treeToArray<K extends Comparable<Object>, V>(node: Node<K, V> | null): Entry<K, V>[] {
465    return foldTree(node, new Entry<K, V>[0], treeFolder);
466}
467