[leetcode #208] Implement Trie (Prefix Tree)

Problem

A trie (pronounced as "try") or prefix tree is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. There are various applications of this data structure, such as autocomplete and spellchecker.

Implement the Trie class:

・ Trie() Initializes the trie object.
・ void insert(String word) Inserts the string word into the trie.
・ boolean search(String word) Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
・ boolean startsWith(String prefix) Returns true if there is a previously inserted string word that has the prefix prefix, and false otherwise.

Example 1:

Input
["Trie", "insert", "search", "search", "startsWith", "insert", "search"]
[[], ["apple"], ["apple"], ["app"], ["app"], ["app"], ["app"]]
Output
[null, null, true, false, true, null, true]

Explanation
Trie trie = new Trie();
trie.insert("apple");
trie.search("apple");   // return True
trie.search("app");     // return False
trie.startsWith("app"); // return True
trie.insert("app");
trie.search("app");     // return True

Constraints:

・ 1 <= word.length, prefix.length <= 2000
・ word and prefix consist only of lowercase English letters.
・ At most 3 * 10⁴ calls in total will be made to insert, search, and startsWith.

Idea

Trie는 예전에 string의 빠른 검색과 관련된 문제를 풀 때 다루어 본 자료구조다. 이전의 기억을 떠올리며 Trie를 만드니 그다지 어렵지 않게 풀 수 있었다.

Trie는 TrieNode가 alphabet의 개수만큼 자식 TrieNode를 가진다는 것이 특징이다.

String을 insert할 때는 alphabet에 해당하는 TrieNode를 찾고 없으면 새로 만든다. String을 처음부터 끝까지 탐색하며 문자를 하나씩 TrieNode로 만들어 자식 노드에 더하는 것이 핵심이다. 이 때 마지막 문자인 경우 end flag를 켜서 해당 문자열이 있다는 것을 표시해둔다.

String을 검색할 때는 Trie 내에서 alphabet의 유무를 확인하며 string이 있는지 검색한다. alphabet이 없다고 확인되면 곧바로 false를 리턴하고, 있는 경우 자식 노드에서 다음 문자를 찾는다. 문자열의 끝에 도달했다면 end flag를 확인하여 해당 문자열이 있는지 유무를 리턴한다.

prefix 검색은 일반 검색과 같지만 end flag 확인없이 prefix의 끝에 도달하면 true를 리턴하면 된다.

Solution

class Trie {
    TrieNode root;

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

    public void insert(String word) {
	TrieNode node = root;
	for (int i=0; i < word.length(); i++) {
	    char currentChar = word.charAt(i);
	    if (!node.contains(currentChar)) {
		node.put(currentChar, new TrieNode());
	    }
	    node = node.get(currentChar);
	}
	node.setEnd();
    }

    public boolean search(String word) {
	TrieNode node = root;
	for (int i=0; i < word.length(); i++) {
	    char currentChar = word.charAt(i);
	    if (!node.contains(currentChar))
	        return false;
	    node = node.get(currentChar);
    	}

	return node.isEnd();
    }

    public boolean startsWith(String prefix) {
	TrieNode node = root;

	for (int i=0; i < prefix.length(); i++) {
	    char currentChar = prefix.charAt(i);
	    if (!node.contains(currentChar))
	        return false;
	    node = node.get(currentChar);
	}

        return true;
    }
}

class TrieNode {
    TrieNode[] children;
    boolean end;

    TrieNode() { children = new TrieNode[26]; }

    boolean contains(char c) {
	return children[c - 'a'] != null;
    }

    TrieNode get(char c) {
	return children[c - 'a'];
    }

    void put(char c, TrieNode node) {
	children[c - 'a'] = node;
    }

    void setEnd() {
	end = true;
    }

    boolean isEnd() {
	return end;
    }
}

Reference

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