Skip to content

Commit

Permalink
lib: add internal PriorityQueue class
Browse files Browse the repository at this point in the history
An efficient JS implementation of a binary heap on top of an array with
worst-case O(log n) runtime for all operations, including arbitrary
item removal (unlike O(n) for most binary heap array implementations).

PR-URL: nodejs#20555
Fixes: nodejs#16105
Reviewed-By: Ben Noordhuis <info@bnoordhuis.nl>
Reviewed-By: Jeremiah Senkpiel <fishrock123@rocketmail.com>
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: James M Snell <jasnell@gmail.com>
Reviewed-By: Ruben Bridgewater <ruben@bridgewater.de>
  • Loading branch information
apapirovski committed May 22, 2018
1 parent a5aad24 commit 6f6f7f7
Show file tree
Hide file tree
Showing 4 changed files with 227 additions and 0 deletions.
18 changes: 18 additions & 0 deletions benchmark/util/priority-queue.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
'use strict';

const common = require('../common');

const bench = common.createBenchmark(main, {
n: [1e6]
}, { flags: ['--expose-internals'] });

function main({ n, type }) {
const PriorityQueue = require('internal/priority_queue');
const queue = new PriorityQueue();
bench.start();
for (var i = 0; i < n; i++)
queue.insert(Math.random() * 1e7 | 0);
for (i = 0; i < n; i++)
queue.shift();
bench.end(n);
}
111 changes: 111 additions & 0 deletions lib/internal/priority_queue.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
'use strict';

const kCompare = Symbol('compare');
const kHeap = Symbol('heap');
const kSetPosition = Symbol('setPosition');
const kSize = Symbol('size');

// The PriorityQueue is a basic implementation of a binary heap that accepts
// a custom sorting function via its constructor. This function is passed
// the two nodes to compare, similar to the native Array#sort. Crucially
// this enables priority queues that are based on a comparison of more than
// just a single criteria.

module.exports = class PriorityQueue {
constructor(comparator, setPosition) {
if (comparator !== undefined)
this[kCompare] = comparator;
if (setPosition !== undefined)
this[kSetPosition] = setPosition;

this[kHeap] = new Array(64);
this[kSize] = 0;
}

[kCompare](a, b) {
return a - b;
}

insert(value) {
const heap = this[kHeap];
let pos = ++this[kSize];

if (heap.length === pos)
heap.length *= 2;

const compare = this[kCompare];
const setPosition = this[kSetPosition];
while (pos > 1) {
const parent = heap[pos / 2 | 0];
if (compare(parent, value) <= 0)
break;
heap[pos] = parent;
if (setPosition !== undefined)
setPosition(parent, pos);
pos = pos / 2 | 0;
}
heap[pos] = value;
if (setPosition !== undefined)
setPosition(value, pos);
}

peek() {
return this[kHeap][1];
}

percolateDown(pos) {
const compare = this[kCompare];
const setPosition = this[kSetPosition];
const heap = this[kHeap];
const size = this[kSize];
const item = heap[pos];

while (pos * 2 <= size) {
let childIndex = pos * 2 + 1;
if (childIndex > size || compare(heap[pos * 2], heap[childIndex]) < 0)
childIndex = pos * 2;
const child = heap[childIndex];
if (compare(item, child) <= 0)
break;
if (setPosition !== undefined)
setPosition(child, pos);
heap[pos] = child;
pos = childIndex;
}
heap[pos] = item;
if (setPosition !== undefined)
setPosition(item, pos);
}

removeAt(pos) {
const heap = this[kHeap];
const size = --this[kSize];
heap[pos] = heap[size + 1];
heap[size + 1] = undefined;

if (size > 0)
this.percolateDown(1);
}

remove(value) {
const heap = this[kHeap];
const pos = heap.indexOf(value);
if (pos < 1)
return false;

this.removeAt(pos);

return true;
}

shift() {
const heap = this[kHeap];
const value = heap[1];
if (value === undefined)
return;

this.removeAt(1);

return value;
}
};
1 change: 1 addition & 0 deletions node.gyp
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@
'lib/internal/safe_globals.js',
'lib/internal/net.js',
'lib/internal/os.js',
'lib/internal/priority_queue.js',
'lib/internal/process/esm_loader.js',
'lib/internal/process/methods.js',
'lib/internal/process/next_tick.js',
Expand Down
97 changes: 97 additions & 0 deletions test/parallel/test-priority-queue.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
// Flags: --expose-internals
'use strict';

require('../common');

const assert = require('assert');
const PriorityQueue = require('internal/priority_queue');

{
// Checks that the queue is fundamentally correct.
const queue = new PriorityQueue();
for (let i = 15; i > 0; i--)
queue.insert(i);

for (let i = 1; i < 16; i++) {
assert.strictEqual(queue.peek(), i);
assert.strictEqual(queue.shift(), i);
}

assert.strictEqual(queue.shift(), undefined);

// Reverse the order.
for (let i = 1; i < 16; i++)
queue.insert(i);

for (let i = 1; i < 16; i++) {
assert.strictEqual(queue.shift(), i);
}

assert.strictEqual(queue.shift(), undefined);
}

{
// Checks that the queue is capable of resizing and fitting more elements.
const queue = new PriorityQueue();
for (let i = 2048; i > 0; i--)
queue.insert(i);

for (let i = 1; i < 2049; i++) {
assert.strictEqual(queue.shift(), i);
}

assert.strictEqual(queue.shift(), undefined);
}

{
// Checks that remove works as expected.
const queue = new PriorityQueue();
for (let i = 16; i > 0; i--)
queue.insert(i);

const removed = [5, 10, 15];
for (const id of removed)
assert(queue.remove(id));

assert(!queue.remove(100));
assert(!queue.remove(-100));

for (let i = 1; i < 17; i++) {
if (removed.indexOf(i) < 0)
assert.strictEqual(queue.shift(), i);
}

assert.strictEqual(queue.shift(), undefined);
}

{
// Make a max heap with a custom sort function.
const queue = new PriorityQueue((a, b) => b - a);
for (let i = 1; i < 17; i++)
queue.insert(i);

for (let i = 16; i > 0; i--) {
assert.strictEqual(queue.shift(), i);
}

assert.strictEqual(queue.shift(), undefined);
}

{
// Make a min heap that accepts objects as values, which necessitates
// a custom sorting function. In addition, add a setPosition function
// as 2nd param which provides a reference to the node in the heap
// and allows speedy deletions.
const queue = new PriorityQueue((a, b) => {
return a.value - b.value;
}, (node, pos) => (node.position = pos));
for (let i = 1; i < 17; i++)
queue.insert({ value: i, position: null });

for (let i = 1; i < 17; i++) {
assert.strictEqual(queue.peek().value, i);
queue.removeAt(queue.peek().position);
}

assert.strictEqual(queue.peek(), undefined);
}

0 comments on commit 6f6f7f7

Please sign in to comment.