트라이

알고리즘

Python impl

"""
https://leetcode.com/problems/implement-trie-prefix-tree
"""


class Trie:
    class Node:
        def __init__(self, letter: str) -> None:
            self.c = letter
            self.e: dict[str, Trie.Node] = {}
            self.end = False

    def __init__(self) -> None:
        self.root = Trie.Node("")

    def insert(self, word: str) -> None:
        """
        1. prefix가 없을때까지 내려간다
        2. 한 글자씩 Node를 생성하고 그 child를 만든다
        3. 마지막 글자인 경우, end를 True로 설정한다
        """
        cur = self.root
        i = 0

        # 1.
        while i < len(word) and cur.e.get(word[i]):
            cur = cur.e[word[i]]
            i += 1

        # 2.
        for c in word[i:]:
            cur.e[c] = Trie.Node(c)
            cur = cur.e[c]

        # 3.
        cur.end = True

    def search(self, word: str) -> bool:
        cur = self.root
        i = 0

        while i < len(word) and cur.e.get(word[i]):
            cur = cur.e[word[i]]
            i += 1

        if i < len(word) or not cur.end:
            return False
        return True

    def startsWith(self, prefix: str) -> bool:
        cur = self.root
        i = 0

        while i < len(prefix) and cur.e.get(prefix[i]):
            cur = cur.e[prefix[i]]
            i += 1

        if i < len(prefix):
            return False
        return True


# Your Trie object will be instantiated and called as such:

word = "apple"
prefix = "app"
obj = Trie()
obj.insert(word)
assert obj.search(word)
assert obj.startsWith(prefix)

C++ impl

unordered_map {{cpp}}를 사용하여 문제풀이

/**
https://leetcode.com/problems/implement-trie-prefix-tree
*/

#include <assert.h>
#include <string>
#include <unordered_map>
#include <utility>

using namespace std;

class Trie {
    struct Node {
        using Dict = unordered_map<char, Node *>;
        char ch;
        Dict dict;
        bool end = false;

        Node(char ch) {
            this->ch = ch;
            this->dict = Dict{};
        }
    };

  private:
    Node *root = nullptr;

  public:
    Trie() { this->root = new Node(0); }

    /**
     * 1. prefix가 없을때까지 내려간다
     * 2. 한 글자씩 Node를 생성하고 그 child를 만든다
     * 3. 마지막 글자인 경우, end를 True로 설정한다
     */
    void insert(string word) {
        auto [cur, i] = this->search_prefix(word);

        // 한 글자씩 Node를 생성하고 child를 만든다
        for (auto iter = word.cbegin() + i; iter != word.cend(); ++iter) {
            const char c = *iter;
            cur->dict[c] = new Node(c);
            cur = cur->dict[c];
        }

        // 마지막 글자에 한해서 end를 true로 설정한다.
        cur->end = true;
    }

    bool search(string word) {
        auto [cur, i] = this->search_prefix(word);
        if (i < word.size() || !cur->end) {
            return false;
        }
        return true;
    }

    bool startsWith(string prefix) {
        auto [cur, i] = this->search_prefix(prefix);
        if (i < prefix.size()) {
            return false;
        }
        return true;
    }

  private:
    /// @brief Trie에 등록돼 있는 word의 마지막 접두어를 찾고 그 인덱스를
    /// 리턴한다
    /// @return [Node *, index]
    pair<Node *, size_t> search_prefix(string word) {
        auto *cur = this->root;
        size_t i = 0;
        size_t wordlen = word.size();

        while (i < wordlen && cur->dict.find(word[i]) != cur->dict.end()) {
            cur = cur->dict[word[i]];
            i++;
        }

        return make_pair(cur, i);
    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * bool param_2 = obj->search(word);
 * bool param_3 = obj->startsWith(prefix);
 */

int main(int argc, char const *argv[]) {
    const string word = "apple";
    const string prefix = "app";
    Trie *obj = new Trie();
    obj->insert(word);
    assert(obj->search(word));
    assert(obj->startsWith(prefix));
}