트라이

문제

• 백준.14725.개미굴
• 백준.5052.전화번호 목록

설명

트라이는 문자열을 효율적으로 저장하고 탐색할 수 있는 K진 트리의 자료구조이다.

원하는 문자열을 찾는 데 들어가는 시간 복잡도는 O(N)이다.

각 단계마다 문자열들을 key로 설정하고 이후 탐색을 진행하는 방식으로 탐색을 최적화한다.

이때 중복된 단계에서 문자열이 어떨때는 마지막 문자열일수도, 아닐수도 있기 때문에 boolean으로 이 단계가 끝인지 여부또한 저장한다.

루트 노드의 경우 key값이 없음에 유의한다.

구현

TrieNode의 구현은 다음과 같다.

// use array
struct TrieNode {
  // 모든 char에 대해 선언했으므로 문제에 주어진 26자에 대해서만 최적화도 가능하다
  struct TrieNode *children[256] = {
      NULL,
  };

  bool isEnd = false;
};

// use hashmap
struct TrieNode {
  unordered_map<char, TrieNode *> children;
  bool isEnd;

  TrieNode() { isEnd = false; }
};

이를 이용해 트라이를 다음과 같이 구현할 수 있다.

class Trie {
 private:
  struct TrieNode {
    unordered_map<char, TrieNode *> children;
    bool isEnd;

    TrieNode() { isEnd = false; }
  };

  TrieNode *root = new TrieNode;

 public:
  void insert(string word) {
    TrieNode *node = root;

    for (char c : word) {
      if (!node->children[c]) {
        node->children[c] = new TrieNode;
      }

      node = node->children[c];
    }

    node->isEnd = true;
  }

  bool search(string word) {
    TrieNode *node = root;

    for (char c : word) {
      if (!node->children[c]) {
        return false;
      }

      node = node->children[c];
    }

    return node && node->isEnd;
  }

  bool startsWith(string prefix) {
    TrieNode *node = root;

    for (char c : prefix) {
      if (!node->children[c]) {
        return false;
      }

      node = node->children[c];
    }

    return node;
  }
};

삭제 가능한 트라이

class Trie {
 private:
  struct TrieNode {
    unordered_map<char, TrieNode *> children;
    bool isEnd;

    TrieNode() { isEnd = false; }
  };

  TrieNode *root = new TrieNode;

  bool isEmpty(TrieNode *cur) {
    if (cur->children.size() > 0) return false;
    return true;
  }

  void erase(string &word, int index, TrieNode *cur) {
    int length = word.length();

    if (index == length) {
      cur->isEnd = false;
      return;
    }

    int key = word[index];

    erase(word, index + 1, cur->children[key]);

    if (isEmpty(cur->children[key])) {
      cur->children[key] = NULL;
      delete cur->children[key];
    }
  }

 public:
  void insert(string word) {
    TrieNode *node = root;

    for (char c : word) {
      if (!node->children[c]) {
        node->children[c] = new TrieNode;
      }

      node = node->children[c];
    }

    node->isEnd = true;
  }

  bool has(string word) {
    TrieNode *node = root;

    for (char c : word) {
      if (!node->children[c]) {
        return false;
      }

      node = node->children[c];
    }

    return node && node->isEnd;
  }

  bool startsWith(string prefix) {
    TrieNode *node = root;

    for (char c : prefix) {
      if (!node->children[c]) {
        return false;
      }

      node = node->children[c];
    }

    return node;
  }

  void erase(string word) { erase(word, 0, root); }
};