Make it possible to change priority of PriorityQueue.

trekhleb · trekhleb · commit 8b057b10d0e9 · 2018-05-03T06:34:32.000+03:00
diff --git a/src/data-structures/heap/MinHeap.js b/src/data-structures/heap/MinHeap.js
@@ -136,16 +136,18 @@ export default class MinHeap {
 
   /**
    * @param {*} item
+   * @param {Comparator} [customFindingComparator]
    * @return {MinHeap}
    */
-  remove(item) {
+  remove(item, customFindingComparator) {
     // Find number of items to remove.
     const numberOfItemsToRemove = this.find(item).length;
+    const customComparator = customFindingComparator || this.compare;
 
     for (let iteration = 0; iteration < numberOfItemsToRemove; iteration += 1) {
       // We need to find item index to remove each time after removal since
       // indices are being change after each heapify process.
-      const indexToRemove = this.find(item).pop();
+      const indexToRemove = this.find(item, customComparator).pop();
 
       // If we need to remove last child in the heap then just remove it.
       // There is no need to heapify the heap afterwards.
@@ -180,13 +182,15 @@ export default class MinHeap {
 
   /**
    * @param {*} item
+   * @param {Comparator} [customComparator]
    * @return {Number[]}
    */
-  find(item) {
+  find(item, customComparator) {
     const foundItemIndices = [];
+    const comparator = customComparator || this.compare;
 
     for (let itemIndex = 0; itemIndex < this.heapContainer.length; itemIndex += 1) {
-      if (this.compare.equal(item, this.heapContainer[itemIndex])) {
+      if (comparator.equal(item, this.heapContainer[itemIndex])) {
         foundItemIndices.push(itemIndex);
       }
     }
diff --git a/src/data-structures/priority-queue/PriorityQueue.js b/src/data-structures/priority-queue/PriorityQueue.js
@@ -10,11 +10,54 @@ export default class PriorityQueue extends MinHeap {
     this.compare = new Comparator(this.comparePriority.bind(this));
   }
 
+  /**
+   * @param {*} item
+   * @param {number} [priority]
+   * @return {PriorityQueue}
+   */
   add(item, priority = 0) {
     this.priorities[item] = priority;
     super.add(item);
+
+    return this;
+  }
+
+  /**
+   * @param {*} item
+   * @param {Comparator} [customFindingComparator]
+   * @return {PriorityQueue}
+   */
+  remove(item, customFindingComparator) {
+    super.remove(item, customFindingComparator);
+    delete this.priorities[item];
+
+    return this;
+  }
+
+  /**
+   * @param {*} item
+   * @param {number} priority
+   * @return {PriorityQueue}
+   */
+  changePriority(item, priority) {
+    const customFindingComparator = new Comparator((a, b) => {
+      if (a === b) {
+        return 0;
+      }
+      return a < b ? -1 : 1;
+    });
+
+    this.remove(item, customFindingComparator);
+    this.add(item, priority);
+
+    return this;
   }
 
+  /**
+   * @param {*} a
+   * @param {*} b
+   * @return {number}
+   */
   comparePriority(a, b) {
     if (this.priorities[a] === this.priorities[b]) {
       return 0;
diff --git a/src/data-structures/priority-queue/__test__/PriorityQueue.test.js b/src/data-structures/priority-queue/__test__/PriorityQueue.test.js
@@ -33,4 +33,58 @@ describe('PriorityQueue', () => {
     expect(priorityQueue.poll()).toBe(10);
     expect(priorityQueue.poll()).toBe(5);
   });
+
+  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.changePriority(100, 10);
+    priorityQueue.changePriority(10, 20);
+
+    expect(priorityQueue.poll()).toBe(200);
+    expect(priorityQueue.poll()).toBe(5);
+    expect(priorityQueue.poll()).toBe(100);
+    expect(priorityQueue.poll()).toBe(10);
+  });
+
+  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.changePriority(200, 10);
+    priorityQueue.changePriority(10, 20);
+
+    expect(priorityQueue.poll()).toBe(100);
+    expect(priorityQueue.poll()).toBe(5);
+    expect(priorityQueue.poll()).toBe(200);
+    expect(priorityQueue.poll()).toBe(10);
+  });
+
+  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.changePriority(200, 10);
+    priorityQueue.changePriority(10, 20);
+
+    priorityQueue.add(15, 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);
+  });
 });
