Skip to content

Commit f5550ae

Browse files
ignlgignlg
committed
refactor: optimize remove and sorting methods in Heap and HeapAsync classes for improved performance
1 parent 5c58a96 commit f5550ae

File tree

2 files changed

+71
-90
lines changed

2 files changed

+71
-90
lines changed

‎src/Heap.ts

+36-41
Original file line numberDiff line numberDiff line change
@@ -614,24 +614,28 @@ export class Heap<T> implements Iterable<T> {
614614
* @return {Boolean} True if the heap was modified
615615
*/
616616
remove(o?: T, callbackFn: IsEqual<T> = Heap.defaultIsEqual): boolean {
617-
if (this.length > 0) {
618-
if (o === undefined) {
619-
this.pop();
620-
return true;
621-
} else {
622-
const idx = this.indexOf(o, callbackFn);
623-
if (idx >= 0) {
624-
if (idx === 0) {
625-
this.pop();
626-
} else if (idx === this.length - 1) {
627-
this.heapArray.pop();
628-
} else {
629-
this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
630-
this._sortNodeUp(idx);
631-
this._sortNodeDown(idx);
632-
}
633-
return true;
617+
if (!this.heapArray.length) return false;
618+
if (o === undefined) {
619+
this.pop();
620+
return true;
621+
}
622+
const queue = [0];
623+
while (queue.length) {
624+
const idx = queue.shift() as number;
625+
if (callbackFn(this.heapArray[idx], o)) {
626+
if (idx === 0) {
627+
this.pop();
628+
} else if (idx === this.heapArray.length - 1) {
629+
this.heapArray.pop();
630+
} else {
631+
this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
632+
this._sortNodeUp(idx);
633+
this._sortNodeDown(idx);
634634
}
635+
return true;
636+
} else if (this.compare(this.heapArray[idx], o) <= 0) {
637+
const children = Heap.getChildrenIndexOf(idx).filter((c) => c < this.heapArray.length);
638+
queue.push(...children);
635639
}
636640
}
637641
return false;
@@ -808,26 +812,20 @@ export class Heap<T> implements Iterable<T> {
808812
* @param {Number} i Index of the node
809813
*/
810814
_sortNodeDown(i: number): void {
811-
let moveIt = i < this.heapArray.length - 1;
812-
const self = this.heapArray[i];
813-
814-
const getPotentialParent = (best: number, j: number) => {
815-
if (this.heapArray.length > j && this.compare(this.heapArray[j], this.heapArray[best]) < 0) {
816-
best = j;
815+
const { length } = this.heapArray;
816+
while (true) {
817+
const left = 2 * i + 1;
818+
const right = left + 1;
819+
let best = i;
820+
if (left < length && this.compare(this.heapArray[left], this.heapArray[best]) < 0) {
821+
best = left;
817822
}
818-
return best;
819-
};
820-
821-
while (moveIt) {
822-
const childrenIdx = Heap.getChildrenIndexOf(i);
823-
const bestChildIndex = childrenIdx.reduce(getPotentialParent, childrenIdx[0]);
824-
const bestChild = this.heapArray[bestChildIndex];
825-
if (typeof bestChild !== 'undefined' && this.compare(self, bestChild) > 0) {
826-
this._moveNode(i, bestChildIndex);
827-
i = bestChildIndex;
828-
} else {
829-
moveIt = false;
823+
if (right < length && this.compare(this.heapArray[right], this.heapArray[best]) < 0) {
824+
best = right;
830825
}
826+
if (best === i) break;
827+
this._moveNode(i, best);
828+
i = best;
831829
}
832830
}
833831

@@ -836,15 +834,12 @@ export class Heap<T> implements Iterable<T> {
836834
* @param {Number} i Index of the node
837835
*/
838836
_sortNodeUp(i: number): void {
839-
let moveIt = i > 0;
840-
while (moveIt) {
837+
while (i > 0) {
841838
const pi = Heap.getParentIndexOf(i);
842-
if (pi >= 0 && this.compare(this.heapArray[pi], this.heapArray[i]) > 0) {
839+
if (pi >= 0 && this.compare(this.heapArray[i], this.heapArray[pi]) < 0) {
843840
this._moveNode(i, pi);
844841
i = pi;
845-
} else {
846-
moveIt = false;
847-
}
842+
} else break;
848843
}
849844
}
850845

‎src/HeapAsync.ts

+35-49
Original file line numberDiff line numberDiff line change
@@ -537,30 +537,28 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
537537
* @return {Boolean} True if the heap was modified
538538
*/
539539
async remove(o?: T, fn: AsyncIsEqual<T> = HeapAsync.defaultIsEqual): Promise<boolean> {
540-
if (this.length > 0) {
541-
if (o === undefined) {
542-
await this.pop();
540+
if (!this.heapArray.length) return false;
541+
if (o === undefined) {
542+
await this.pop();
543+
return true;
544+
}
545+
const queue = [0];
546+
while (queue.length) {
547+
const idx = queue.shift() as number;
548+
if (await fn(this.heapArray[idx], o)) {
549+
if (idx === 0) {
550+
await this.pop();
551+
} else if (idx === this.heapArray.length - 1) {
552+
this.heapArray.pop();
553+
} else {
554+
this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
555+
await this._sortNodeUp(idx);
556+
await this._sortNodeDown(idx);
557+
}
543558
return true;
544559
} else {
545-
let idx = -1;
546-
for (let i = 0; i < this.heapArray.length; ++i) {
547-
if (await fn(this.heapArray[i], o)) {
548-
idx = i;
549-
break;
550-
}
551-
}
552-
if (idx >= 0) {
553-
if (idx === 0) {
554-
await this.pop();
555-
} else if (idx === this.length - 1) {
556-
this.heapArray.pop();
557-
} else {
558-
this.heapArray.splice(idx, 1, this.heapArray.pop() as T);
559-
await this._sortNodeUp(idx);
560-
await this._sortNodeDown(idx);
561-
}
562-
return true;
563-
}
560+
const children = HeapAsync.getChildrenIndexOf(idx).filter((c) => c < this.heapArray.length);
561+
queue.push(...children);
564562
}
565563
}
566564
return false;
@@ -737,29 +735,20 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
737735
* @param {Number} i Index of the node
738736
*/
739737
async _sortNodeDown(i: number): Promise<void> {
740-
let moveIt = i < this.heapArray.length - 1;
741-
const self = this.heapArray[i];
742-
743-
const getPotentialParent = async (best: number, j: number) => {
744-
if (this.heapArray.length > j && (await this.compare(this.heapArray[j], this.heapArray[best])) < 0) {
745-
best = j;
738+
const length = this.heapArray.length;
739+
while (true) {
740+
const left = 2 * i + 1;
741+
const right = left + 1;
742+
let best = i;
743+
if (left < length && (await this.compare(this.heapArray[left], this.heapArray[best])) < 0) {
744+
best = left;
746745
}
747-
return best;
748-
};
749-
750-
while (moveIt) {
751-
const childrenIdx = HeapAsync.getChildrenIndexOf(i);
752-
let bestChildIndex = childrenIdx[0];
753-
for (let j = 1; j < childrenIdx.length; ++j) {
754-
bestChildIndex = await getPotentialParent(bestChildIndex, childrenIdx[j]);
755-
}
756-
const bestChild = this.heapArray[bestChildIndex];
757-
if (typeof bestChild !== 'undefined' && (await this.compare(self, bestChild)) > 0) {
758-
this._moveNode(i, bestChildIndex);
759-
i = bestChildIndex;
760-
} else {
761-
moveIt = false;
746+
if (right < length && (await this.compare(this.heapArray[right], this.heapArray[best])) < 0) {
747+
best = right;
762748
}
749+
if (best === i) break;
750+
this._moveNode(i, best);
751+
i = best;
763752
}
764753
}
765754

@@ -768,15 +757,12 @@ export class HeapAsync<T> implements Iterable<Promise<T>> {
768757
* @param {Number} i Index of the node
769758
*/
770759
async _sortNodeUp(i: number): Promise<void> {
771-
let moveIt = i > 0;
772-
while (moveIt) {
760+
while (i > 0) {
773761
const pi = HeapAsync.getParentIndexOf(i);
774-
if (pi >= 0 && (await this.compare(this.heapArray[pi], this.heapArray[i])) > 0) {
762+
if (pi >= 0 && (await this.compare(this.heapArray[i], this.heapArray[pi])) < 0) {
775763
this._moveNode(i, pi);
776764
i = pi;
777-
} else {
778-
moveIt = false;
779-
}
765+
} else break;
780766
}
781767
}
782768

0 commit comments

Comments
 (0)