208. Implement Trie (Prefix Tree) - cocoder39/coco39_LC GitHub Wiki

class TrieNode:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.childres = collections.defaultdict(TrieNode)
        self.isWord = False

class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TrieNode()
        

    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        cur = self.root
        for ch in word:
            cur = cur.childres[ch]
        cur.isWord = True
        
    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        cur = self.root
        for ch in word:
            if ch not in cur.childres:
                return False
            cur = cur.childres[ch]
        return cur.isWord
        
    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        cur = self.root
        for ch in prefix:
            if ch not in cur.childres:
                return False
            cur = cur.childres[ch]
        return True

class TrieNode {
public:
    // Initialize your data structure here.
    TrieNode() {
        isWord = false;
        children = vector<TrieNode*>(26, nullptr);
    }
    bool isWord; //end of a word
    vector<TrieNode*> children;
};

class Trie {
public:
    Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie.
    void insert(string word) {
        if (word.empty())   return;
        TrieNode* p = root;
        for (auto c : word) {
            int idx = c - 'a';
            if (! (p->children)[idx])   (p->children)[idx] = new TrieNode();
            p = p->children[idx];
        }
        p->isWord = true;
    }

    // Returns if the word is in the trie.
    bool search(string word) {
        TrieNode* p = find(word);
        return p && p->isWord;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    bool startsWith(string prefix) {
        return find(prefix) != nullptr;
    }

private:
    TrieNode* root;
    
    TrieNode* find(string& word) {
        if (word.empty())   return nullptr;
        TrieNode* p = root;
        for (auto c : word) {
            int idx = c - 'a';
            if (! (p->children)[idx])  return nullptr;
            p = (p->children)[idx];
        }
        return p;
    }
};