Trie 자료구조 JavaScript

Trie...

Trie 자료구조 라는게 있다. 문자열 검색을 빠르게 해주는 자료구조이다. 요즘 문자열에 관한 코딩 문제들이 정말 자주 나오는데, 그 때마다 오픈채팅방 사람들이 Trie로 풀었다, KMP로 풀었다. 말이 많길래 이번 기회에 알아보았다.

Trie에 대한 내용은 검색을 하면 더 자세하게 나온다.
Trie는 문자열들을 하나하나 쪼개어 tree 구조에 넣음으로써 검색을 더 빠르게 한다.
M의 길이를 가지는 N개의 문자열에 대해서 부분문자열의 경우를 찾는 경우, 2중 for문을 통해 할 경우 O(M * N^2)의 시간복잡도를 가진다. 이렇게 되면 많은 문제에서 TLE를 맛볼 수 있다.

trie는 한개를 탐색할 때 O(M), 전체 O(MN)으로 풀 수가 있다. 물론 Trie를 생성할 때 O(MN)이 걸리지만 O(M * N^2) 보다 훨씬 짧다.
N의 개수가 늘어날 수록 훨씬 효과적일 것이다.

javascript로 어떻게 구현하면 되는지를 더 중점적으로 다룰려고 한다.

코드

const Node = function () { // 각각의 노드
  this.keys = new Map(); // 현재 노드에서 갈 수 있는 문자들을 저장할 Map
  this.end = false; // 해당 문자에서 끝나는 문자열이 있는지 확인.
  this.count = 0; // 지나 갔을 때 1 씩 증가시켜, 1보다 클 경우 지나간 문자가 자신말고 더 있다는 것을 알 수 있다.
  this.plusCount = function () {
    this.count += 1;
  };
  this.setEnd = function () {
    this.end = true;
  };
  this.isEnd = function () {
    return this.end;
  };
};

const Trie = function () { // Trie 자료구조
  this.root = new Node();

  this.add = function (input, node = this.root) { // Trie에 문자열을 추가 하는 메소드.
    if (input.length === 0) { // 끝날 경우 end로 표시하고 종료
      node.setEnd();
      return;
    } else if (!node.keys.has(input[0])) { // 첫번째 문자가 없으면 추가 해줌.
      node.keys.set(input[0], new Node());
    }
    node.plusCount(); // 지나갔다고 1 증가 시킨다.
    return this.add(input.substr(1), node.keys.get(input[0])); // 첫번재 문자를 뺀 문자열과, 해당 값의 Map으로 재귀
  };

  this.isSubWord = function (word) { // 자신이 부분 문자열이 되는지 확인.
    let node = this.root;
    while (word.length > 0) {
      if (!node.keys.has(word[0])) {
        return false;
      } else {
        node = node.keys.get(word[0]);
        word = word.substr(1);
      }
    }
    return node.count > 0 ? true : false; // 끝까지 갔을 때 자신의 count가 0보다 크면 부분 문자열이다. 마지막 문자열일 때는 +1을 해주지 않으므로 마지막까지 갔을 때 node의 count가 0이면 그 문자는 유니크한 것.
  };

  this.isWord = function (word) { // 있는 문자인지 확인 하는 메소드
    let node = this.root;
    while (word.length > 1) {
      if (!node.keys.has(word[0])) {
        return false;
      } else {
        node = node.keys.get(word[0]);
        word = word.substr(1);
      }
    } // 마지막 문자가 해당 key에 있고 그 node가 end이면 true 아닐 경우 false
    return node.keys.has(word) && node.keys.get(word).isEnd() ? true : false;
  };
 
  this.print = function () { // 모든 문자열을 array로 return 해주는 함수.
    let words = [];
    let search = function (node = this.root, string) {
      if (node.keys.size !== 0) {
        for (let letter of node.keys.keys()) {
          search(node.keys.get(letter), string.concat(letter)); // search 재귀, concat으로 현재 문자열은 그대로 두고 letter를 더한 문자열을 반환하게 한다.
        }
        if (node.isEnd()) {
          words.push(string);
        }
      } else {
        string.length > 0 ? words.push(string) : undefined;
        return;
      }
    };
    search(this.root, '');
    return words.length > 0 ? words : null;
  };
};

const myTrie = new Trie();
myTrie.add('ball');
myTrie.add('bat');
myTrie.add('doll');
myTrie.add('dark');
myTrie.add('do');
myTrie.add('sense');
console.log('doll', myTrie.isWord('doll'));
console.log('dor', myTrie.isWord('dor'));
console.log('do', myTrie.isSubWord('do'));
console.log('doll', myTrie.isSubWord('doll'));
console.log('zzz', myTrie.isSubWord('zzz'));
console.log(myTrie.print());

/* 
doll true
dor false
do true
doll false
zzz false
[ 'ball', 'bat', 'doll', 'do', 'dark', 'sense' ]
*/

참고

https://www.youtube.com/watch?v=7XmS8McW_1U&t=200s