diff --git a/src/data-structures/heap/MinHeap.js b/src/data-structures/heap/MinHeap.js index bf23f415..5f32f53f 100644 --- a/src/data-structures/heap/MinHeap.js +++ b/src/data-structures/heap/MinHeap.js @@ -141,8 +141,8 @@ export default class MinHeap { */ remove(item, customFindingComparator) { // Find number of items to remove. - const numberOfItemsToRemove = this.find(item).length; const customComparator = customFindingComparator || this.compare; + const numberOfItemsToRemove = this.find(item, customComparator).length; for (let iteration = 0; iteration < numberOfItemsToRemove; iteration += 1) { // We need to find item index to remove each time after removal since diff --git a/src/data-structures/heap/__test__/MinHeap.test.js b/src/data-structures/heap/__test__/MinHeap.test.js index f4ce4ac3..b2b8a7b1 100644 --- a/src/data-structures/heap/__test__/MinHeap.test.js +++ b/src/data-structures/heap/__test__/MinHeap.test.js @@ -1,4 +1,5 @@ import MinHeap from '../MinHeap'; +import Comparator from '../../../utils/comparator/Comparator'; describe('MinHeap', () => { it('should create an empty min heap', () => { @@ -147,4 +148,25 @@ describe('MinHeap', () => { expect(minHeap.remove(3).toString()).toEqual('4'); expect(minHeap.remove(4).toString()).toEqual(''); }); + + it('should be possible to remove items from heap with custom finding comparator', () => { + const minHeap = new MinHeap(); + minHeap.add('dddd'); + minHeap.add('ccc'); + minHeap.add('bb'); + minHeap.add('a'); + + expect(minHeap.toString()).toBe('a,bb,ccc,dddd'); + + const comparator = new Comparator((a, b) => { + if (a.length === b.length) { + return 0; + } + + return a.length < b.length ? -1 : 1; + }); + + minHeap.remove('hey', comparator); + expect(minHeap.toString()).toBe('a,bb,dddd'); + }); });