Implement Trie (Prefix Tree)

Published by

on

Problem

https://leetcode.com/explore/learn/card/trie/147/basic-operations/1047/

Solution

/**
1. Initialize: cur = root
2. for each char c in target string S:
3.      if cur does not have a child c:
4.          cur.children[c] = new Trie node
5.      cur = cur.children[c]
6. cur is the node which represents the string S
**/
class Trie {
    
    class TrieNode {
        private Map<Character, TrieNode> childrenMap = new HashMap<>();
        private boolean isWord = false;
        
        public Map<Character, TrieNode> getChildrenMap() {
            return this.childrenMap;
        }
        
        public boolean isWord() {
            return this.isWord;
        }
        
        public void setWord() {
            this.isWord = true;
        }
    }

    TrieNode root;
    
    public Trie() {
        root = new TrieNode();
    }
    
    public void insert(String word) {
        char[] charArr = word.toCharArray();
        TrieNode current = root;
        for (int i=0;i<charArr.length;i++) {
            char c = charArr[i];
            if (!current.getChildrenMap().containsKey(c)) {
                TrieNode next = new TrieNode();
                current.getChildrenMap().put(c, next);
            }
            current = current.getChildrenMap().get(c);
        }
        current.setWord();
    }
    
    public boolean search(String word) {
        char[] charArr = word.toCharArray();
        TrieNode current = root;
        for (int i=0;i<charArr.length;i++) {
            char c = charArr[i];
            if (!current.getChildrenMap().containsKey(c)) {
                return false;
            }
            current = current.getChildrenMap().get(c);
        }
        
        if (!current.isWord()) {
            return false;
        }
        
        return true;
    }
    
    public boolean startsWith(String prefix) {
        char[] charArr = prefix.toCharArray();
        TrieNode current = root;
        for (int i=0;i<charArr.length;i++) {
            char c = charArr[i];
            if (!current.getChildrenMap().containsKey(c)) {
                return false;
            }
            
            current = current.getChildrenMap().get(c);
        }
        
        return true;
    }
    
}

/**
 * Your Trie object will be instantiated and called as such:
 * Trie obj = new Trie();
 * obj.insert(word);
 * boolean param_2 = obj.search(word);
 * boolean param_3 = obj.startsWith(prefix);
 */
  • Trie 문제의 대표적인 코드
  • TrieNode 라는 클래스를 선언하는 것이 포인트!
  • 시작할때 root 를 만들고, current 라는 포인터를 기준으로 순회하는 것이 중요.

Leave a comment