feat(util): 🎸 add AvlSet implementation

Author: streamich

src/util/trees/avl/AvlSet.ts

@@ -0,0 +1,104 @@
+import {insert, insertLeft, remove, insertRight, print} from './util';
+import {printTree} from '../../print/printTree';
+import {findOrNextLower, first, next} from '../util';
+import type {Printable} from '../../print/types';
+import type {Comparator, HeadlessNode} from '../types';
+import type {AvlNodeReference, IAvlTreeNode} from './types';
+
+export class AvlSetNode<V> implements IAvlTreeNode<V, void> {
+  public p: AvlSetNode<V> | undefined = undefined;
+  public l: AvlSetNode<V> | undefined = undefined;
+  public r: AvlSetNode<V> | undefined = undefined;
+  public bf: number = 0;
+  public v: undefined = undefined;
+  constructor(public readonly k: V) {}
+}
+
+const defaultComparator = (a: unknown, b: unknown) => (a === b ? 0 : (a as any) < (b as any) ? -1 : 1);
+
+export class AvlSet<V> implements Printable {
+  public root: AvlSetNode<V> | undefined = undefined;
+  public readonly comparator: Comparator<V>;
+
+  constructor(comparator?: Comparator<V>) {
+    this.comparator = comparator || defaultComparator;
+  }
+
+  private insert(value: V): AvlNodeReference<AvlSetNode<V>> {
+    const item = new AvlSetNode<V>(value);
+    this.root = insert(this.root, item, this.comparator);
+    return item;
+  }
+
+  public add(value: V): AvlNodeReference<AvlSetNode<V>> {
+    const root = this.root;
+    if (!root) return this.insert(value);
+    const comparator = this.comparator;
+    let next: AvlSetNode<V> | undefined = root,
+      curr: AvlSetNode<V> | undefined = next;
+    let cmp: number = 0;
+    do {
+      curr = next;
+      cmp = comparator(value, curr.k);
+      if (cmp === 0) return curr;
+    } while ((next = cmp < 0 ? (curr.l as AvlSetNode<V>) : (curr.r as AvlSetNode<V>)));
+    const node = new AvlSetNode<V>(value);
+    this.root =
+      cmp < 0 ? (insertLeft(root, node, curr) as AvlSetNode<V>) : (insertRight(root, node, curr) as AvlSetNode<V>);
+    return node;
+  }
+
+  private find(k: V): AvlNodeReference<AvlSetNode<V>> | undefined {
+    const comparator = this.comparator;
+    let curr: AvlSetNode<V> | undefined = this.root;
+    while (curr) {
+      const cmp = comparator(k, curr.k);
+      if (cmp === 0) return curr;
+      curr = cmp < 0 ? (curr.l as AvlSetNode<V>) : (curr.r as AvlSetNode<V>);
+    }
+    return undefined;
+  }
+
+  public del(k: V): void {
+    const node = this.find(k);
+    if (!node) return;
+    this.root = remove(this.root, node as AvlSetNode<V>);
+  }
+
+  public clear(): void {
+    this.root = undefined;
+  }
+
+  public has(k: V): boolean {
+    return !!this.find(k);
+  }
+
+  public size(): number {
+    const root = this.root;
+    if (!root) return 0;
+    let curr = first(root);
+    let size = 1;
+    while ((curr = next(curr as HeadlessNode) as AvlSetNode<V> | undefined)) size++;
+    return size;
+  }
+
+  public isEmpty(): boolean {
+    return !this.root;
+  }
+
+  public getOrNextLower(k: V): AvlSetNode<V> | undefined {
+    return (findOrNextLower(this.root, k, this.comparator) as AvlSetNode<V>) || undefined;
+  }
+
+  public forEach(fn: (node: AvlSetNode<V>) => void): void {
+    const root = this.root;
+    if (!root) return;
+    let curr = first(root);
+    do fn(curr!);
+    while ((curr = next(curr as HeadlessNode) as AvlSetNode<V> | undefined));
+  }
+
+  public toString(tab: string): string {
+    return this.constructor.name + printTree(tab, [(tab) => print(this.root, tab)]);
+  }
+}

src/util/trees/avl/__tests__/AvlSet.spec.ts

@@ -0,0 +1,61 @@
+import {AvlSet} from '../AvlSet';
+
+test('can add numbers to set', () => {
+  const set = new AvlSet<number>();
+  expect(set.size()).toBe(0);
+  expect(set.has(1)).toBe(false);
+  set.add(1);
+  expect(set.size()).toBe(1);
+  expect(set.has(1)).toBe(true);
+  set.add(24);
+  set.add(42);
+  set.add(42);
+  expect(set.size()).toBe(3);
+  expect(set.has(24)).toBe(true);
+  expect(set.has(42)).toBe(true);
+  expect(set.has(25)).toBe(false);
+});
+
+test('can remove numbers from set', () => {
+  const set = new AvlSet<number>();
+  set.add(1);
+  set.add(24);
+  set.add(42);
+  expect(set.has(1)).toBe(true);
+  expect(set.has(24)).toBe(true);
+  expect(set.has(42)).toBe(true);
+  set.del(24);
+  expect(set.has(1)).toBe(true);
+  expect(set.has(24)).toBe(false);
+  expect(set.has(42)).toBe(true);
+  set.del(1);
+  expect(set.has(1)).toBe(false);
+  expect(set.has(24)).toBe(false);
+  expect(set.has(42)).toBe(true);
+  expect(set.size()).toBe(1);
+  set.del(42);
+  expect(set.has(1)).toBe(false);
+  expect(set.has(24)).toBe(false);
+  expect(set.has(42)).toBe(false);
+  expect(set.size()).toBe(0);
+});
+
+test('can store structs', () => {
+  class Struct { constructor (public x: number, public y: number) {} }
+  const set = new AvlSet<Struct>((a, b) => {
+    const dx = a.x - b.x;
+    return dx === 0 ? a.y - b.y : dx;
+  });
+  set.add(new Struct(0, 0));
+  set.add(new Struct(0, 1));
+  expect(set.size()).toBe(2);
+  set.del(new Struct(0, 0));
+  expect(set.size()).toBe(1);
+  expect(set.has(new Struct(0, 0))).toBe(false);
+  expect(set.has(new Struct(0, 1))).toBe(true);
+  set.add(new Struct(2, 3));
+  set.add(new Struct(3, 3));
+  expect(set.size()).toBe(3);
+  expect(set.has(new Struct(3, 3))).toBe(true);
+  expect(set.has(new Struct(2, 3))).toBe(true);
+});