Add support for complex data to PriorityQueue

src/data-structures/priority-queue/PriorityQueue.js

@@ -4,87 +4,112 @@ import Comparator from '../../utils/comparator/Comparator';
 // It is the same as min heap except that when comparing to elements
 // we take into account not element's value but rather its priority.
 export default class PriorityQueue extends MinHeap {
-  constructor() {
+  /**
+   * @constructor
+   * @param {function|undefined} compareValueFunction
+   */
+  constructor(compareValueFunction) {
     super();
-    this.priorities = {};
+    // Map data structure supports using any value for key type
+    // e.g. functions, objects, or primitives
+    this.priorities = new Map();
     this.compare = new Comparator(this.comparePriority.bind(this));
+    this.compareValue = new Comparator(compareValueFunction);
   }
 
   /**
    * @param {*} item
-   * @param {number} [priority]
+   * @param {number} priority
    * @return {PriorityQueue}
    */
   add(item, priority = 0) {
-    this.priorities[item] = priority;
+    this.priorities.set(item, priority);
     super.add(item);
 
     return this;
   }
 
   /**
    * @param {*} item
-   * @param {Comparator} [customFindingComparator]
+   * @param {Comparator|function|undefined} maybeComparator
    * @return {PriorityQueue}
    */
-  remove(item, customFindingComparator) {
-    super.remove(item, customFindingComparator);
-    delete this.priorities[item];
+  remove(item, maybeComparator) {
+    const comparator = this.getValueComparator(maybeComparator);
+    super.remove(item, comparator);
+    this.priorities.delete(item);
 
     return this;
   }
 
   /**
    * @param {*} item
    * @param {number} priority
+   * @param {Comparator|function|undefined} maybeComparator
    * @return {PriorityQueue}
    */
-  changePriority(item, priority) {
-    this.remove(item, new Comparator(this.compareValue));
+  changePriority(item, priority, maybeComparator) {
+    const comparator = this.getValueComparator(maybeComparator);
+    this.remove(item, comparator);
     this.add(item, priority);
 
     return this;
   }
 
   /**
    * @param {*} item
+   * @param {Comparator|function|undefined} maybeComparator
    * @return {Number[]}
    */
-  findByValue(item) {
-    return this.find(item, new Comparator(this.compareValue));
+  findByValue(item, maybeComparator) {
+    const comparator = this.getValueComparator(maybeComparator);
+    return this.find(item, comparator);
   }
 
   /**
    * @param {*} item
+   * @param {Comparator|function|undefined} maybeComparator
    * @return {boolean}
    */
-  hasValue(item) {
-    return this.findByValue(item).length > 0;
+  hasValue(item, maybeComparator) {
+    const comparator = this.getValueComparator(maybeComparator);
+    return this.findByValue(item, comparator).length > 0;
   }
 
   /**
-   * @param {*} a
-   * @param {*} b
-   * @return {number}
+   * @param {Comparator|function|undefined} maybeComparator
+   * @return {Comparator}
    */
-  comparePriority(a, b) {
-    if (this.priorities[a] === this.priorities[b]) {
-      return 0;
+  getValueComparator(maybeComparator) {
+    if (maybeComparator == null) {
+      return this.compareValue;
+    }
+
+    if (maybeComparator instanceof Comparator) {
+      return maybeComparator;
     }
 
-    return this.priorities[a] < this.priorities[b] ? -1 : 1;
+    if (maybeComparator instanceof Function) {
+      return new Comparator(maybeComparator);
+    }
+
+    throw new TypeError(
+      'Invalid comparator type\n'
+      + 'Must be one of: Comparator | Function | undefined\n'
+      + `Given: ${typeof maybeComparator}`,
+    );
   }
 
   /**
    * @param {*} a
    * @param {*} b
    * @return {number}
    */
-  compareValue(a, b) {
-    if (a === b) {
+  comparePriority(a, b) {
+    if (this.priorities.get(a) === this.priorities.get(b)) {
       return 0;
     }
 
-    return a < b ? -1 : 1;
+    return this.priorities.get(a) < this.priorities.get(b) ? -1 : 1;
   }
 }

src/data-structures/priority-queue/__test__/PriorityQueue.test.js

@@ -1,5 +1,11 @@
 import PriorityQueue from '../PriorityQueue';
 
+const JOB1 = { type: 'job1' };
+const JOB2 = { type: 'job2' };
+const JOB3 = { type: 'job3' };
+const JOB4 = { type: 'job4' };
+const JOB5 = { type: 'job5' };
+
 describe('PriorityQueue', () => {
   it('should create default priority queue', () => {
     const priorityQueue = new PriorityQueue();
@@ -10,94 +16,118 @@ describe('PriorityQueue', () => {
   it('should insert items to the queue and respect priorities', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    expect(priorityQueue.peek()).toBe(10);
+    priorityQueue.add(JOB1, 1);
+    expect(priorityQueue.peek()).toBe(JOB1);
 
-    priorityQueue.add(5, 2);
-    expect(priorityQueue.peek()).toBe(10);
+    priorityQueue.add(JOB2, 2);
+    expect(priorityQueue.peek()).toBe(JOB1);
 
-    priorityQueue.add(100, 0);
-    expect(priorityQueue.peek()).toBe(100);
+    priorityQueue.add(JOB3, 0);
+    expect(priorityQueue.peek()).toBe(JOB3);
   });
 
   it('should poll from queue with respect to priorities', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    priorityQueue.add(5, 2);
-    priorityQueue.add(100, 0);
-    priorityQueue.add(200, 0);
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+    priorityQueue.add(JOB4, 0);
 
-    expect(priorityQueue.poll()).toBe(100);
-    expect(priorityQueue.poll()).toBe(200);
-    expect(priorityQueue.poll()).toBe(10);
-    expect(priorityQueue.poll()).toBe(5);
+    expect(priorityQueue.poll()).toBe(JOB3);
+    expect(priorityQueue.poll()).toBe(JOB4);
+    expect(priorityQueue.poll()).toBe(JOB1);
+    expect(priorityQueue.poll()).toBe(JOB2);
   });
 
   it('should be possible to change priority of internal nodes', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    priorityQueue.add(5, 2);
-    priorityQueue.add(100, 0);
-    priorityQueue.add(200, 0);
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+    priorityQueue.add(JOB4, 0);
 
-    priorityQueue.changePriority(100, 10);
-    priorityQueue.changePriority(10, 20);
+    priorityQueue.changePriority(JOB4, 10);
+    priorityQueue.changePriority(JOB1, 20);
 
-    expect(priorityQueue.poll()).toBe(200);
-    expect(priorityQueue.poll()).toBe(5);
-    expect(priorityQueue.poll()).toBe(100);
-    expect(priorityQueue.poll()).toBe(10);
+    expect(priorityQueue.poll()).toBe(JOB3);
+    expect(priorityQueue.poll()).toBe(JOB2);
+    expect(priorityQueue.poll()).toBe(JOB4);
+    expect(priorityQueue.poll()).toBe(JOB1);
   });
 
   it('should be possible to change priority of head node', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    priorityQueue.add(5, 2);
-    priorityQueue.add(100, 0);
-    priorityQueue.add(200, 0);
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+    priorityQueue.add(JOB4, 0);
 
-    priorityQueue.changePriority(200, 10);
-    priorityQueue.changePriority(10, 20);
+    priorityQueue.changePriority(JOB3, 10);
+    priorityQueue.changePriority(JOB1, 20);
 
-    expect(priorityQueue.poll()).toBe(100);
-    expect(priorityQueue.poll()).toBe(5);
-    expect(priorityQueue.poll()).toBe(200);
-    expect(priorityQueue.poll()).toBe(10);
+    expect(priorityQueue.poll()).toBe(JOB4);
+    expect(priorityQueue.poll()).toBe(JOB2);
+    expect(priorityQueue.poll()).toBe(JOB3);
+    expect(priorityQueue.poll()).toBe(JOB1);
   });
 
   it('should be possible to change priority along with node addition', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    priorityQueue.add(5, 2);
-    priorityQueue.add(100, 0);
-    priorityQueue.add(200, 0);
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+    priorityQueue.add(JOB4, 0);
 
-    priorityQueue.changePriority(200, 10);
-    priorityQueue.changePriority(10, 20);
+    priorityQueue.changePriority(JOB4, 10);
+    priorityQueue.changePriority(JOB1, 20);
 
-    priorityQueue.add(15, 15);
+    priorityQueue.add(JOB5, 15);
 
-    expect(priorityQueue.poll()).toBe(100);
-    expect(priorityQueue.poll()).toBe(5);
-    expect(priorityQueue.poll()).toBe(200);
-    expect(priorityQueue.poll()).toBe(15);
-    expect(priorityQueue.poll()).toBe(10);
+    expect(priorityQueue.poll()).toBe(JOB3);
+    expect(priorityQueue.poll()).toBe(JOB2);
+    expect(priorityQueue.poll()).toBe(JOB4);
+    expect(priorityQueue.poll()).toBe(JOB5);
+    expect(priorityQueue.poll()).toBe(JOB1);
   });
 
   it('should be possible to search in priority queue by value', () => {
     const priorityQueue = new PriorityQueue();
 
-    priorityQueue.add(10, 1);
-    priorityQueue.add(5, 2);
-    priorityQueue.add(100, 0);
-    priorityQueue.add(200, 0);
-    priorityQueue.add(15, 15);
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+    priorityQueue.add(JOB4, 0);
+    priorityQueue.add(JOB5, 15);
+
+    const job6 = { type: 'job6' };
+
+    expect(priorityQueue.hasValue(job6)).toBe(false);
+    expect(priorityQueue.hasValue(JOB5)).toBe(true);
+  });
+
+  it('should accept a custom compareValue function', () => {
+    const compareByType = (a, b) => {
+      if (a.type === b.type) {
+        return 0;
+      }
+
+      return a.type < b.type ? -1 : 1;
+    };
+
+    const priorityQueue = new PriorityQueue(compareByType);
+
+    priorityQueue.add(JOB1, 1);
+    priorityQueue.add(JOB2, 2);
+    priorityQueue.add(JOB3, 0);
+
+    const existingJobType = { type: 'job1' };
+    const newJobType = { type: 'job4' };
 
-    expect(priorityQueue.hasValue(70)).toBe(false);
-    expect(priorityQueue.hasValue(15)).toBe(true);
+    expect(priorityQueue.hasValue(existingJobType)).toBe(true);
+    expect(priorityQueue.hasValue(newJobType)).toBe(false);
   });
 });