Skip to content

Instantly share code, notes, and snippets.

@callemo
Created March 25, 2019 11:16
Show Gist options
  • Save callemo/c14bad4cba817d969a1bee01e1b7e7f8 to your computer and use it in GitHub Desktop.
Save callemo/c14bad4cba817d969a1bee01e1b7e7f8 to your computer and use it in GitHub Desktop.
/* trie.js */
class TrieNode {
constructor(value) {
this.value = value;
this.children = new Map();
this.isEndOfWord = false;
}
}
class Trie {
constructor() {
this.size = 0;
this.root = new TrieNode(null);
}
has(word) {
const path = this._findPathToNode(word);
if (!path.length) {
return false;
}
return path[path.length - 1].isEndOfWord;
}
add(word) {
if (!word) {
return;
}
let currentNode = this.root;
let nextNode;
for (const ch of word) {
nextNode = currentNode.children.get(ch);
if (!nextNode) {
nextNode = new TrieNode(ch);
currentNode.children.set(ch, nextNode);
}
currentNode = nextNode;
}
currentNode.isEndOfWord = true;
this.size += 1;
}
delete(word) {
const path = this._findPathToNode(word);
if (!path.length) {
return false;
}
let currentNode = path[path.length - 1];
if (!currentNode.isEndOfWord) {
return false;
}
currentNode.isEndOfWord = false;
this.size -= 1;
if (currentNode.children.size) {
return true;
}
let parentNode;
for (let i = path.length - 2; i >= 0; i--) {
if (currentNode.isEndOfWord || currentNode.children.size > 1) {
break;
}
parentNode = path[i];
parentNode.children.delete(currentNode.value);
currentNode = parentNode;
}
return true;
}
search(prefix) {
const path = this._findPathToNode(prefix);
if (!path.length) {
return [];
}
const results = [];
let currentNode = path[path.length - 1];
if (currentNode.isEndOfWord) {
results.push(prefix);
}
/* dfs search */
const stack = [];
let currentWord = prefix;
for (const childNode of Array.from(currentNode.children.values()).reverse()) {
stack.push([currentWord, childNode]);
}
while (stack.length) {
[currentWord, currentNode] = stack.pop();
currentWord += currentNode.value;
if (currentNode.isEndOfWord) {
results.push(currentWord);
}
for (const childNode of Array.from(currentNode.children.values()).reverse()) {
stack.push([currentWord, childNode]);
}
}
return results;
}
/* Aho-Corasick string matching */
match() {
}
clear() {
this.size = 0;
this.root = new TrieNode(null);
}
print() {
let depth, currentNode;
const stack = [[0, this.root]];
while (stack.length) {
[depth, currentNode] = stack.pop();
for (const childNode of Array.from(currentNode.children.values()).reverse()) {
stack.push([depth + 1, childNode]);
}
/* eslint no-console:off */
console.log(
" ".repeat(depth) +
(depth ? currentNode.value : "*") +
(currentNode.isEndOfWord ? "." : "")
);
}
}
_findPathToNode(word) {
if (!word) {
return [];
}
const path = [];
let currentNode = this.root;
let nextNode;
path.push(currentNode);
for (const ch of word) {
nextNode = currentNode.children.get(ch);
if (!nextNode) {
return [];
}
currentNode = nextNode;
path.push(currentNode);
}
return path;
}
}
module.exports = { Trie };
/* trie.test.js */
const assert = require("assert");
const { Trie } = require("./trie");
describe("Trie", function() {
let t;
const words = ["sin", "sing", "the", "then", "thin", "this", "tin"];
beforeEach(function() {
t = new Trie();
});
describe(".add", function() {
it("inserts a word", function() {
assert.equal(t.size, 0);
t.add("example");
assert.equal(t.size, 1);
assert(t.has("example") === true);
});
});
describe(".search", function() {
it("retrieves matching words", function() {
for (const word of words) {
t.add(word);
}
assert.equal(t.size, words.length);
assert.deepEqual(t.search(), []);
assert.deepEqual(t.search("notaword"), []);
assert.deepEqual(t.search("sing"), ["sing"]);
assert.deepEqual(t.search("sin"), ["sin", "sing"]);
assert.deepEqual(t.search("t"), ["the", "then", "thin", "this", "tin"]);
});
});
describe(".delete", function() {
it("removes a word", function() {
for (const word of words) {
t.add(word);
}
assert.equal(t.delete("then"), true);
assert.equal(t.size, words.length - 1);
assert(!t.has("then"));
assert(t.has("thin"));
assert.deepEqual(t.search("t"), ["the", "thin", "this", "tin"]);
});
it("removes a word root", function() {
t.add("bat");
t.add("bath");
t.add("bathroom");
t.add("bathtub");
assert(t.delete("bath"));
assert.deepEqual(t.search("b"), ["bat", "bathroom", "bathtub"]);
});
});
describe(".clear", function() {
it("removes all words", function() {
for (const word of words) {
t.add(word);
}
t.clear();
assert.equal(t.size, 0);
assert.equal(t.root.children.size, 0);
assert.deepEqual(t.search("t"), []);
});
});
});
/* eslint-env mocha */
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment