Implement Trie (Prefix Tree) — Insert/Search/StartsWith

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Implement a Trie with:

  • insert(word) — insert a word
  • search(word) — returns true if word exists
  • startsWith(prefix) — returns true if any word starts with prefix

Solutions

C++ — HashMap children

class Trie {
    struct Node {
        unordered_map<char,Node*> ch;
        bool end=false;
    };
    Node* root=new Node();
public:
    void insert(string w){
        auto* c=root;
        for(char x:w){if(!c->ch[x])c->ch[x]=new Node();c=c->ch[x];}
        c->end=true;
    }
    bool search(string w){
        auto* c=root;
        for(char x:w){if(!c->ch.count(x))return false;c=c->ch[x];}
        return c->end;
    }
    bool startsWith(string p){
        auto* c=root;
        for(char x:p){if(!c->ch.count(x))return false;c=c->ch[x];}
        return true;
    }
};

Java — array children

class Trie {
    private Trie[] ch=new Trie[26];
    private boolean end;
    public void insert(String w){
        Trie c=this;
        for(char x:w.toCharArray()){
            int i=x-'a';
            if(c.ch[i]==null)c.ch[i]=new Trie();
            c=c.ch[i];
        }
        c.end=true;
    }
    public boolean search(String w){
        Trie c=this;
        for(char x:w.toCharArray()){int i=x-'a';if(c.ch[i]==null)return false;c=c.ch[i];}
        return c.end;
    }
    public boolean startsWith(String p){
        Trie c=this;
        for(char x:p.toCharArray()){int i=x-'a';if(c.ch[i]==null)return false;c=c.ch[i];}
        return true;
    }
}

Python

class TrieNode:
    def __init__(self):
        self.children={}
        self.is_end=False

class Trie:
    def __init__(self):
        self.root=TrieNode()
    def insert(self, word: str) -> None:
        node=self.root
        for c in word:
            if c not in node.children: node.children[c]=TrieNode()
            node=node.children[c]
        node.is_end=True
    def search(self, word: str) -> bool:
        node=self.root
        for c in word:
            if c not in node.children: return False
            node=node.children[c]
        return node.is_end
    def startsWith(self, prefix: str) -> bool:
        node=self.root
        for c in prefix:
            if c not in node.children: return False
            node=node.children[c]
        return True

Complexity

  • Time: O(L) per operation, L = word length
  • Space: O(total characters * 26) array, or O(total characters) hashmap

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro