feat: add type overloads for heap

harttle · Jun 14, 2021 · d2057b6 · d2057b6
1 parent e3c7c5e
commit d2057b6
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/heap.ts b/heap.ts
@@ -1,7 +1,16 @@
-export class Heap<T> {
+export class Heap<T=number> {
   data: Array<T | null>
   lt: (i: number, j: number) => boolean
-  constructor (data: T[] = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
+  constructor ()
+  constructor (data: T[])
+  constructor (cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: T[], cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)), cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)) = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
+    if (typeof data === 'function') {
+      cmp = data
+      data = []
+    }
     this.data = [null, ...data]
     this.lt = (i, j) => cmp(this.data[i]!, this.data[j]!)
     for (let i = this.size(); i > 0; i--) this.heapify(i)
@@ -47,7 +56,12 @@ export class RemovableHeap<T> {
   heap: Heap<T>
   counts: Map<T, number>
   _invalidCount: number
-  constructor (data: T[] = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
+  constructor ()
+  constructor (data: T[])
+  constructor (cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: T[], cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)), cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)) = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
     this.heap = new Heap<T>(data, cmp)
     this.counts = new Map()
     this._invalidCount = 0
@@ -98,7 +112,12 @@ export class RemovableHeap<T> {
 export class DoubleRemovableHeap<T> {
   min: RemovableHeap<T>
   max: RemovableHeap<T>
-  constructor (data: T[] = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
+  constructor ()
+  constructor (data: T[])
+  constructor (cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: T[], cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)), cmp: (lhs: T, rhs: T) => boolean)
+  constructor (data: (T[] | ((lhs: T, rhs: T) => boolean)) = [], cmp = (lhs: T, rhs: T) => lhs < rhs) {
     this.min = new RemovableHeap(data, cmp)
     this.max = new RemovableHeap(data, (lhs, rhs) => !cmp(lhs, rhs))
   }

diff --git a/test/heap.spec.ts b/test/heap.spec.ts
@@ -41,7 +41,7 @@ describe('heap', () => {
       expect(ans).toEqual('001122334455')
     })
     it('should support custom less', () => {
-      const heap = new Heap<number>([], (l, r) => l > r)
+      const heap = new Heap<number>((l, r) => l > r)
       for (let i = 0; i <= 5; i++) heap.push(i)
       let ans = ''
       for (let i = 0; i <= 5; i++) {