Add Trie.deleteWord and TrieNode.removeChild (trekhleb#181)

appleJax · trekhleb · commit d25eff49e60d · 2018-08-27T15:33:16.000+03:00
diff --git a/src/data-structures/trie/Trie.js b/src/data-structures/trie/Trie.js
@@ -24,6 +24,35 @@ export default class Trie {
     return this;
   }
 
+  /**
+   * @param {string} word
+   * @return {Trie}
+   */
+  deleteWord(word) {
+    function depthFirstDelete(currentNode, charIndex) {
+      if (charIndex >= word.length) return;
+
+      const character = word[charIndex];
+      const nextNode = currentNode.getChild(character);
+
+      if (nextNode == null) return;
+
+      depthFirstDelete(nextNode, charIndex + 1);
+
+      if (charIndex === word.length - 1) {
+        nextNode.isCompleteWord = false;
+      }
+
+      // childNode is deleted only if:
+      // - childNode has NO children
+      // - childNode.isCompleteWord === false
+      currentNode.removeChild(character);
+    }
+
+    depthFirstDelete(this.head, 0);
+    return this;
+  }
+
   /**
    * @param {string} word
    * @return {string[]}
diff --git a/src/data-structures/trie/TrieNode.js b/src/data-structures/trie/TrieNode.js
@@ -37,6 +37,31 @@ export default class TrieNode {
     return childNode;
   }
 
+  /**
+   * @param {string} character
+   * @return {TrieNode}
+   */
+  removeChild(character) {
+    function isSafeToDelete(node) {
+      return (
+        node
+        && !node.isCompleteWord
+        && node.children.getKeys().length === 0
+      );
+    }
+
+    const childNode = this.getChild(character);
+
+    // delete childNode only if:
+    // - childNode has NO children
+    // - childNode.isCompleteWord === false
+    if (isSafeToDelete(childNode)) {
+      this.children.delete(character);
+    }
+
+    return this;
+  }
+
   /**
    * @param {string} character
    * @return {boolean}
diff --git a/src/data-structures/trie/__test__/Trie.test.js b/src/data-structures/trie/__test__/Trie.test.js
@@ -23,6 +23,18 @@ describe('Trie', () => {
     expect(trie.head.getChild('c').getChild('a').getChild('t').toString()).toBe('t*');
   });
 
+  it('should delete words from trie', () => {
+    const trie = new Trie();
+
+    trie.addWord('carpet');
+    trie.addWord('car');
+    expect(trie.doesWordExist('carpet')).toBe(true);
+
+    trie.deleteWord('carpet');
+    expect(trie.doesWordExist('carpet')).toEqual(false);
+    expect(trie.doesWordExist('car')).toEqual(true);
+  });
+
   it('should suggests next characters', () => {
     const trie = new Trie();
 
diff --git a/src/data-structures/trie/__test__/TrieNode.test.js b/src/data-structures/trie/__test__/TrieNode.test.js
@@ -18,6 +18,36 @@ describe('TrieNode', () => {
     expect(trieNode.toString()).toBe('c:a,o');
   });
 
+  describe('removing child nodes', () => {
+    it('should delete child node if the child node has NO children', () => {
+      const trieNode = new TrieNode('c');
+      trieNode.addChild('a');
+      expect(trieNode.hasChild('a')).toBe(true);
+
+      trieNode.removeChild('a');
+      expect(trieNode.hasChild('a')).toBe(false);
+    });
+
+    it('should NOT delete child node if the child node has children', () => {
+      const trieNode = new TrieNode('c');
+      trieNode.addChild('a');
+      const childNode = trieNode.getChild('a');
+      childNode.addChild('r');
+
+      trieNode.removeChild('a');
+      expect(trieNode.hasChild('a')).toEqual(true);
+    });
+
+    it('should NOT delete child node if the child node completes a word', () => {
+      const trieNode = new TrieNode('c');
+      const IS_COMPLETE_WORD = true;
+      trieNode.addChild('a', IS_COMPLETE_WORD);
+
+      trieNode.removeChild('a');
+      expect(trieNode.hasChild('a')).toEqual(true);
+    });
+  });
+
   it('should get child nodes', () => {
     const trieNode = new TrieNode('c');
 
