[leetcode #211] Design Add and Search Words Data Structure

Problem

Design a data structure that supports adding new words and finding if a string matches any previously added string.

Implement the WordDictionary class:

・ WordDictionary() Initializes the object.
・ void addWord(word) Adds word to the data structure, it can be matched later.
・ bool search(word) Returns true if there is any string in the data structure that matches word or false otherwise. word may contain dots '.' where dots can be matched with any letter.

Example:

・ Input
["WordDictionary","addWord","addWord","addWord","search","search","search","search"]
[[],["bad"],["dad"],["mad"],["pad"],["bad"],[".ad"],["b.."]]
・ Output
[null,null,null,null,false,true,true,true]

・ Explanation
WordDictionary wordDictionary = new WordDictionary();
wordDictionary.addWord("bad");
wordDictionary.addWord("dad");
wordDictionary.addWord("mad");
wordDictionary.search("pad"); // return False
wordDictionary.search("bad"); // return True
wordDictionary.search(".ad"); // return True
wordDictionary.search("b.."); // return True

Constraints:

・ 1 <= word.length <= 500
・ word in addWord consists lower-case English letters.
・ word in search consist of  '.' or lower-case English letters.
・ At most 50000 calls will be made to addWord and search.

Idea

주어진 단어들을 저장하고 특정 단어가 존재하는 지 확인하는 클래스를 설계하는 문제다.

Brute-force로 풀면 매번 저장되어 있는 단어를 일일이 비교하기 때문에 Time Limit Exceeded가 뜰 수 밖에 없다. (최대 50,000번의 함수가 호출된며, 만약 모든 String의 길이가 500이라면 25,000,000의 비교를 하게 된다.)

TrieNode를 쓰면 매번 일일이 단어를 찾으면서 비교를 할 필요가 없다. TrieNode는 각 char를 Node로 만들어 Tree를 구성하는 것으로 단어 검색에 유용한 자료구조다.

WordDictionary 생성자에서 TrieNode를 하나 만든다.

addWord에서는 word의 각 문자를 TrieNode로 만들어 더한다. 가장 먼저 등장한 문자가 상위 노드로 존재하게 된다. 문자열이 끝나면 끝을 표시하는 문자 하나를 TrieNode로 만들어 더한다. 이는 문자열을 찾을 때 해당 문자열의 마지막 문자가 존재하는지 확인하기 위함이다.

search에서는 주어진 문자열을 찾는데 재귀함수를 이용하면 편하다. 재귀함수에서는 문자열이 비었고 마지막 문자를 표시하는 Node가 존재하면 해당 단어를 찾은 걸로 판단해 true를 리턴한다. 문자열을 찾는 도중에 '.'가 나왔을 경우 해당 노드의 자식 노드들에 한해 다음 substring을 전부 찾도록 재귀함수를 호출한다. 이는 '.'가 모든 문자에 해당될 수 있기 때문이다. '.'가 아닌 경우는 문자에 해당하는 노드를 찾고 존재할 경우 다음 substring에 대해 재귀함수를 호출하고, 노드가 없을 경우 곧바로 false를 리턴한다.

나는 알파벳 문자 26개에 해당하는 array를 가진 TrieNode 대신 ArrayList를 사용해 매번 리스트 전체를 탐색하게 했다. 이렇게 하면 똑같이 O(n)이 되겠지만 최대 26배가 더 걸리므로 당연히 정석적인 TrieNode를 사용하는 걸 추천한다.

Solution

class WordDictionary {
    private Node root;

    public WordDictionary() {
        root = new Node(' ');
    }

    public void addWord(String word) {
        Node node = root;
        for (int i=0; i < word.length(); i++) {
            node = node.add(word.charAt(i));
        }
        node.add('\n');
    }

    public boolean search(String word) {
        return searchNode(root, word);
    }

    private boolean searchNode(Node node, String word) {
        boolean res = false;

        if (word.isEmpty()) {
            for (Node child: node.children) {
                if (child.c == '\n') {
                    return true;
                }
            }
            return false;
        }

        char c = word.charAt(0);

        if (c == '.') {
            for (Node child : node.children) {
                res |= searchNode(child, word.substring(1));
            }

            return res;
        }

        for (Node child : node.children) {
            if (child.c == c) {
                return searchNode(child, word.substring(1));
            }
        }

        return false;
    }

    class Node {
        char c;
        List<Node> children;

        Node(char c) {
            this.c = c;
            this.children = new ArrayList<>();
        }

        Node add(char c) {
            for (Node node : children) {
                if (node.c == c) {
                    return node;
                }
            }
            Node newNode = new Node(c);
            this.children.add(newNode);

            return newNode;
        }
    }
}

/**
 * Your WordDictionary object will be instantiated and called as such:
 * WordDictionary obj = new WordDictionary();
 * obj.addWord(word);
 * boolean param_2 = obj.search(word);
 */

TrieNode를 쓰면 훨씬 효율적으로 풀 수 있다.

public class WordDictionary {
    public class TrieNode {
        public TrieNode[] children = new TrieNode[26];
        public String item = "";
    }

    private TrieNode root = new TrieNode();

    public void addWord(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            if (node.children[c - 'a'] == null) {
                node.children[c - 'a'] = new TrieNode();
            }
            node = node.children[c - 'a'];
        }
        node.item = word;
    }

    public boolean search(String word) {
        return match(word.toCharArray(), 0, root);
    }

    private boolean match(char[] chs, int k, TrieNode node) {
        if (k == chs.length) return !node.item.equals("");   
        if (chs[k] != '.') {
            return node.children[chs[k] - 'a'] != null && match(chs, k + 1, node.children[chs[k] - 'a']);
        } else {
            for (int i = 0; i < node.children.length; i++) {
                if (node.children[i] != null) {
                    if (match(chs, k + 1, node.children[i])) {
                        return true;
                    }
                }
            }
        }
        return false;
    }
}

Reference

https://leetcode.com/problems/design-add-and-search-words-data-structure/