diff --git a/src/util/trees/avl/AvlSet.ts b/src/util/trees/avl/AvlSet.ts new file mode 100644 index 0000000000..e3b66394dc --- /dev/null +++ b/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 implements IAvlTreeNode { + public p: AvlSetNode | undefined = undefined; + public l: AvlSetNode | undefined = undefined; + public r: AvlSetNode | 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 implements Printable { + public root: AvlSetNode | undefined = undefined; + public readonly comparator: Comparator; + + constructor(comparator?: Comparator) { + this.comparator = comparator || defaultComparator; + } + + private insert(value: V): AvlNodeReference> { + const item = new AvlSetNode(value); + this.root = insert(this.root, item, this.comparator); + return item; + } + + public add(value: V): AvlNodeReference> { + const root = this.root; + if (!root) return this.insert(value); + const comparator = this.comparator; + let next: AvlSetNode | undefined = root, + curr: AvlSetNode | 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) : (curr.r as AvlSetNode))); + const node = new AvlSetNode(value); + this.root = + cmp < 0 ? (insertLeft(root, node, curr) as AvlSetNode) : (insertRight(root, node, curr) as AvlSetNode); + return node; + } + + private find(k: V): AvlNodeReference> | undefined { + const comparator = this.comparator; + let curr: AvlSetNode | undefined = this.root; + while (curr) { + const cmp = comparator(k, curr.k); + if (cmp === 0) return curr; + curr = cmp < 0 ? (curr.l as AvlSetNode) : (curr.r as AvlSetNode); + } + return undefined; + } + + public del(k: V): void { + const node = this.find(k); + if (!node) return; + this.root = remove(this.root, node as AvlSetNode); + } + + 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 | undefined)) size++; + return size; + } + + public isEmpty(): boolean { + return !this.root; + } + + public getOrNextLower(k: V): AvlSetNode | undefined { + return (findOrNextLower(this.root, k, this.comparator) as AvlSetNode) || undefined; + } + + public forEach(fn: (node: AvlSetNode) => void): void { + const root = this.root; + if (!root) return; + let curr = first(root); + do fn(curr!); + while ((curr = next(curr as HeadlessNode) as AvlSetNode | undefined)); + } + + public toString(tab: string): string { + return this.constructor.name + printTree(tab, [(tab) => print(this.root, tab)]); + } +} diff --git a/src/util/trees/avl/__tests__/AvlBstMap.spec.ts b/src/util/trees/avl/__tests__/AvlMap.spec.ts similarity index 100% rename from src/util/trees/avl/__tests__/AvlBstMap.spec.ts rename to src/util/trees/avl/__tests__/AvlMap.spec.ts diff --git a/src/util/trees/avl/__tests__/AvlSet.spec.ts b/src/util/trees/avl/__tests__/AvlSet.spec.ts new file mode 100644 index 0000000000..ddf7dddb6a --- /dev/null +++ b/src/util/trees/avl/__tests__/AvlSet.spec.ts @@ -0,0 +1,66 @@ +import {AvlSet} from '../AvlSet'; + +test('can add numbers to set', () => { + const set = new AvlSet(); + 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(); + 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((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); +});