[JavaScript] Bad word finder

bad-words.js

const assert = require("assert");

function generateTrie(array) {
  const trie = { "": {} };
  for (const word of array) {
    let currentNode = trie[""];
    Array.from(word).forEach((char) => {
      if (!currentNode[char]) {
        currentNode[char] = {};
      }
      currentNode = currentNode[char];
    });
  }
  return trie;
}

function hasContainBadWord(input, trie) {
  let currentNode = trie[""];
  for (const char of input) {
    // We do not have that char as badWord starter.
    if (!currentNode[char]) return false;
    // If the current bad word trie ends,
    if (Object.keys(currentNode[char]).length === 0) {
      return true;
    }
    currentNode = currentNode[char];
  }
  return false;
}

// -----------Test cases-----------

// Generator tests,
assert.deepStrictEqual(generateTrie([]), { "": {} });
assert.deepStrictEqual(generateTrie(["a", "b", "c"]), {
  "": { a: {}, b: {}, c: {} },
});
assert.deepStrictEqual(generateTrie(["aaa", "aab", "aba"]), {
  "": {
    a: {
      a: {
        a: {},
        b: {},
      },
      b: {
        a: {},
      },
    },
  },
});

// hasContainBadWord tests,
const badWords = ["amzn", "amazon", "cagatay"]; // move to the constructor class.
const trie = generateTrie(badWords);

assert.deepStrictEqual(hasContainBadWord("amz", trie), false);
assert.deepStrictEqual(hasContainBadWord("amzn", trie), true);
assert.deepStrictEqual(hasContainBadWord("amzns", trie), true);
assert.deepStrictEqual(hasContainBadWord("cc", trie), false);

fixed.js

const assert = require("assert");

// Prefix tree
class Node {
  constructor(char) {
    this.char = char;
    this.isWord = false;
    this.children = {};
  }
}

class Trie {
  constructor() {
    this.root = new Node('');
  }
  insert(word) {
    let current = this.root;
    for (const char of word) {
      if (!current.children[char]) {
        current.children[char] = new Node(char);
      }
      current = current.children[char];
    }
    current.isWord = true;
  }
  search(word, idx) {
    let current = this.root;
    for (let i = idx; i < word.length; i++) {
      const char = word[i];
      if (!current.children[char]) return false;
      if (current.children[char].isWord) return true;
      current = current.children[char];
    }
    return false;
  }
}


function main(input, badWords) {
  const prefixTree = new Trie();
  badWords.forEach(word => prefixTree.insert(word));
  for (let i = 0; i < input.length; i++) {
    if (prefixTree.search(input, i)) {
      return true;
    }
  }
  return false;
}

assert.deepStrictEqual(main("am", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amz", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amzn", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amzns", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("cc", ['amzn', 'am']), false);

optimized.js

const assert = require("assert");

// Prefix tree
class Node {
  constructor(char) {
    this.char = char;
    this.isWord = false;
    this.children = {};
  }
}

class Trie {
  constructor() {
    this.root = new Node('');
  }
  insert(word) {
    let current = this.root;
    for (const char of word) {
      if (!current.children[char]) {
        current.children[char] = new Node(char);
      }
      current = current.children[char];
    }
    current.isWord = true;
  }
  search(word, idx) {
    let current = this.root;
    for (let i = idx; i < word.length; i++) {
      const char = word[i];
      if (!current.children[char]) return false;
      if (current.children[char].isWord) return true;
      current = current.children[char];
    }
    return false;
  }
}


function main(input, badWords) {
  const prefixTree = new Trie();
  const uniqInput = new Set(input);
  badWords.forEach(word => {
    if (!uniqInput.has(word[0]) || word.length > input.length) {
      // Pass.
    } else {
      prefixTree.insert(word);
    }
  });
  for (let i = 0; i < input.length; i++) {
    if (prefixTree.search(input, i)) {
      return true;
    }
  }
  return false;
}

assert.deepStrictEqual(main("am", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amz", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amzn", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("amzns", ['amzn', 'am', 'cc']), true);
assert.deepStrictEqual(main("cc", ['amzn', 'am']), false);