본문 바로가기
알고리즘 문제풀이/leetcode

[leetcode 208] Implement Trie (Prefix Tree)

by m2162003 2020. 11. 4.

Implement a trie with insert, search, and startsWith methods.

Example:

Trie trie = new Trie(); trie.insert("apple"); trie.search("apple"); // returns true trie.search("app"); // returns false trie.startsWith("app"); // returns true trie.insert("app"); trie.search("app"); // returns true

Note:

  • You may assume that all inputs are consist of lowercase letters a-z.
  • All inputs are guaranteed to be non-empty strings.

문제 풀이:

트라이 구현하기

class TrieNode{
  public:
    TrieNode* child[26];
    bool isEnd;
    
    TrieNode(){
        for(int i=0; i<26; i++){
            child[i] = NULL;
        }
        isEnd = false;
    }
};

class Trie {
    
private:
    TrieNode* root;
    
public:
    /** Initialize your data structure here. */
    
    Trie() {
       root = new TrieNode(); 
    }
    
    /** Inserts a word into the trie. */
    void insert(string word) {
        TrieNode* tmp = root;

        for(int i=0; i<word.length(); i++){
            int idx = word[i] - 'a';
            if(tmp -> child[idx]==0){
                tmp -> child[idx] = new TrieNode();
            }
            tmp = tmp -> child[idx];
        }
        tmp->isEnd = true;
        
    }
    
    /** Returns if the word is in the trie. */
    bool search(string word) {
        TrieNode* tmp = root;
        for(int i=0; i<word.length(); i++){
            int idx = word[i] - 'a';
            if(tmp->child[idx] == NULL) return false;
            tmp = tmp -> child[idx];
        }
        return tmp->isEnd;
    }
    
    /** Returns if there is any word in the trie that starts with the given prefix. */
    bool startsWith(string prefix) {
        TrieNode* tmp = root;
        for(int i=0; i<prefix.length(); i++){
            int idx = prefix[i] - 'a';
            if(tmp->child[idx] == NULL) return false;
            tmp = tmp -> child[idx];
        }
        return true;
    }
};

/**
 * 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);
 */

댓글