diff --git a/src/data-structures/trie/Trie.js b/src/data-structures/trie/Trie.js index a36bd2ef..bd618d20 100644 --- a/src/data-structures/trie/Trie.js +++ b/src/data-structures/trie/Trie.js @@ -29,17 +29,25 @@ export default class Trie { * @return {Trie} */ deleteWord(word) { - function depthFirstDelete(currentNode, charIndex) { - if (charIndex >= word.length) return; + const depthFirstDelete = (currentNode, charIndex = 0) => { + if (charIndex >= word.length) { + // Return if we're trying to delete the character that is out of word's scope. + return; + } const character = word[charIndex]; const nextNode = currentNode.getChild(character); - if (nextNode == null) return; + if (nextNode == null) { + // Return if we're trying to delete a word that has not been added to the Trie. + return; + } + // Go deeper. depthFirstDelete(nextNode, charIndex + 1); - if (charIndex === word.length - 1) { + // Since we're going to delete a word let's un-mark its last character isCompleteWord flag. + if (charIndex === (word.length - 1)) { nextNode.isCompleteWord = false; } @@ -47,9 +55,11 @@ export default class Trie { // - childNode has NO children // - childNode.isCompleteWord === false currentNode.removeChild(character); - } + }; + + // Start depth-first deletion from the head node. + depthFirstDelete(this.head); - depthFirstDelete(this.head, 0); return this; } diff --git a/src/data-structures/trie/__test__/Trie.test.js b/src/data-structures/trie/__test__/Trie.test.js index 92b32f41..4d60ad9c 100644 --- a/src/data-structures/trie/__test__/Trie.test.js +++ b/src/data-structures/trie/__test__/Trie.test.js @@ -35,6 +35,13 @@ describe('Trie', () => { expect(trie.doesWordExist('cart')).toBe(true); expect(trie.doesWordExist('cat')).toBe(true); + // Try to delete not-existing word first. + trie.deleteWord('carpool'); + expect(trie.doesWordExist('carpet')).toBe(true); + expect(trie.doesWordExist('car')).toBe(true); + expect(trie.doesWordExist('cart')).toBe(true); + expect(trie.doesWordExist('cat')).toBe(true); + trie.deleteWord('carpet'); expect(trie.doesWordExist('carpet')).toEqual(false); expect(trie.doesWordExist('car')).toEqual(true);