트라이 (Trie/Prefix Tree) - 2

 

목차

 

추가 메서드의 전체 코드는 다음과 같다.

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

class Main {

    static class Node {
        Map<Character, Node> childNode;
        boolean endOfWord;

        public Node() {
            this.childNode = new HashMap<>();
            this.endOfWord = false;
        }
    }

    static class Trie {
        Node root;

        Trie(Node node) {
            this.root = node;
        }

        void insert(String s) {
            Node current = root;

            for (int i = 0; i < s.length(); i++) {
                current = current.childNode.computeIfAbsent(s.charAt(i), key -> new Node());
            }

            current.endOfWord = true;
        }

        boolean search(String s) {
            Node current = root;

            for (char c : s.toCharArray()) {
                current = current.childNode.getOrDefault(c, null);
                if (current == null) return false;
            }

            return current.endOfWord;
        }

        boolean delete(String s) {
            return delete(root, s, 0);
        }

        private boolean delete(Node current, String s, int next) {

            char childVal = s.charAt(next);
            Node childNode = current.childNode.get(childVal);
            next++;

            if (childNode == null) return false;

            if (next == s.length()) {
                if (!childNode.endOfWord) return false;
                childNode.endOfWord = false;
                if (childNode.childNode.isEmpty()) current.childNode.remove(childVal);

                return true;
            }
            
            if (delete(childNode, s, next)) {
                if (!childNode.endOfWord && childNode.childNode.isEmpty())
                    current.childNode.remove(childVal);
                return true;
            }

            return false;
        }

        List<String> listAll() {
            List<String> results = new ArrayList<>();
            listAll(root, "", results);
            return results;
        }

        private void listAll(Node current, String prefix, List<String> results) {
            if (current.endOfWord) {
                results.add(prefix);
            }

            for (Map.Entry<Character, Node> entry : current.childNode.entrySet()) {
                listAll(entry.getValue(), prefix + entry.getKey(), results);
            }
        }

        private Node findNode(String prefix) {
            Node current = root;
            for (int i = 0; i < prefix.length(); i++) {
                current = current.childNode.getOrDefault(prefix.charAt(i), null);
                if (current == null) return null;
            }
            return current;
        }

        List<String> listWithPrefix(String prefix) {
            List<String> results = new ArrayList<>();

            Node current = findNode(prefix);
            if (current == null) return null;

            listAll(current, prefix, results);
            return results;
        }
    }

    public static void main(String[] args) {

        Node root = new Node();
        Trie trie = new Trie(root);

        trie.insert("node");
        trie.insert("nod");
        trie.insert("mod");

        System.out.println(trie.search("node")); // true
        System.out.println(trie.search("nod")); // true
        System.out.println(trie.listAll()); // [mod, nod, node]
        System.out.println(trie.listWithPrefix("no")); // [nod, node]

        System.out.println(trie.delete("no")); // false
        System.out.println(trie.delete("nodea")); // false
        System.out.println(trie.listAll()); // [mod, nod, node]

        System.out.println(trie.delete("node")); // true
        System.out.println(trie.listAll()); // [mod, nod]

        System.out.println(trie.delete("nod")); // true
        System.out.println(trie.listAll()); // [mod]
    }
}

 

 

5. listAll()의 구현

void listAll(Node current, String prefix, List<String> results) {
    if (current.endOfWord) {
        results.add(prefix);
    }

    for (Map.Entry<Character, Node> entry : current.childNode.entrySet()) {
        listAll(entry.getValue(), prefix + entry.getKey(), results);
    }
}

 

최종적으로는 results를 반환한다.

최근 노드가 endOfWord라면, 이는 곧 독립적인 하나의 단어임을 의미한다. results에 추가한다.

 

아니라면, entrySet을 순회한다.

map 형태인 childNode는 current 노드의 모든 자식 노드를 의미한다.

prefix에 자식 노드의 값을 누적하고, current를 자식 노드로 갱신하며 재귀를 반복한다.

 

유의할 점은, 재귀의 횟수는 자식노드 개수만큼 깊어지나,

results는 List타입이며 레퍼런스 참조이므로, 하나의 인스턴스를 공유한다.

 

 

6. listWithPrefix()의 구현

키워드 자동완성에 쓰이는 메서드이다.

listAll()과의 차이점은, 진입점을 root가 아닌 prefix의 endOfWord에 해당하는 노드로 놓아야 한다.

 

이는 이전에 구현했던 search() 메서드를 boolean이 아닌 Node를 반환하게 하여 응용할 수 있다.

Node findNode(String prefix) {
    Node current = root;
    for (int i = 0; i < prefix.length(); i++) {
        current = current.childNode.getOrDefault(prefix.charAt(i), null);
        if (current == null) return null;
    }
    return current;
}

 

 

listWithPrefix() 또한 listAll()을 응용하여 쉽게 구현할 수 있다.

파라미터에 진입점을 root대신 findNode(prefix)로 설정하고, prefix를 추가해 준다.

 

List<String> listWithPrefix (String prefix) {
    List<String> results = new ArrayList<>();

    Node current = findNode(prefix);
    if(current == null) return null;

    listAll(current, prefix, results);
    return results;
}

 

 

7. delete()의 구현

char childVal = s.charAt(next);
Node childNode = current.childNode.get(childVal);
next++;

if (childNode == null) return false;

 

다른 글에선 next를 idx로 표현하나, 좀 더 명확히 하기위해 변경했다.

처음 진입시, current는 root이며, next는 0이다. s.charAt(0)은 root의 자식 노드이다.

 

이는 delete 하기 위해 있어야 할 자식 노드를 뜻하며, 없을 수도 있다.

없다면 false를 반환해 해당 문자열이 트라이에 존재하지 않으며, delete 동작이 실패했음을 알린다.

 

존재한다면, s.charAt(next)의 값을 찾아 자식 노드로 할당하고, next++ 한다. 

 

 

if (next == s.length()) {
    if (!childNode.endOfWord) return false;
    childNode.endOfWord = false;
    if (childNode.childNode.isEmpty()) current.childNode.remove(childVal);

    return true;
}

next가 s. length() 라는 것은, 자식 노드는 s.length()-1의 값을 가짐을 의미한다. 즉, 삭제하려는 문자열의 마지막이다.

마지막의 단어를 가리키는 노드는 endOfWord가 true여야 한다. 아니라면 해당 단어를 저장하지 않은 것이다. 탐색이 실패했음을 알린다.

 

저장한 단어라면, endOfWord를 false로 변경한다. 이후, 자식이 없다면 삭제한다. 왜 일까?

 

"node"와 "nod"가 저장되어 있을 때, "nod"의 delete 동작을 가정하자.

n-o-d(eow)-e(eow)

이 경우, d는 "nod"의 endOfWord면서, 동시에 "node"의 경로이다.
해당 노드를 삭제할 경우 "node"의 표현을 할 수 없기 때문에 삭제하지 않는다.
endOfWord만 false로 변경한다. 이는 곧 자식 노드가 존재할 때의 경우이다.

최종적으로 다음과 같아진다.
n-o-d-e(eow)


"nod"만 저장되어 있을 때, "nod"의 delete 동작을 가정하자.

n-o-d(eow)

d는 자식 노드가 없다. 이 경우는 삭제할 수 있다.
부모 노드 또한 자식 노드가 없으므로 삭제한다.
최종적으로 모든 노드가 삭제된다.

 

---

 

if (delete(childNode, s, next + 1)) {
    if (!childNode.endOfWord && childNode.childNode.isEmpty())
    	current.childNode.remove(childVal);
    return true;
}
return false;

 

 

재귀적으로 모든 노드를 탐색하며,

마지막 노드가 아닐 경우엔, 다른 단어의 끝이 아니고 자식이 없다면 삭제한다.

delete시 다른 단어에 영향을 미치지 않으려는 동작이라는 점에서 이전의 경우와 비슷한 이유이다.

 

"node"와 "nod"가 저장되어 있을 때, "node"의 delete 동작을 가정하자.

n-o-d(eow)-e(eow)

e를 eow = false로 변경하고, 자식 노드가 없으므로 삭제한다.
eow인 d까지 삭제한다면, "nod"도 같이 삭제가 되어버린다. 따라서 삭제하지 않는다.

최종적으로 다음과 같아진다.
n-o-d(eow)

---

if (delete(childNode, s, next)) {
    if (..) ..;
    return true;
}

return false;

 

또한 주의할 점은, 삭제가 가능함과 상관없이 true를 반환한다.

이는 자식 노드를 삭제할 수 있음을 가리키는 것이 아니고, 탐색을 지속할 수 있음을 뜻한다.

 

문자열에 포함되는 모든 노드를 삭제하지 않더라도,

조건을 만족하는 일부 노드만 삭제하거나 끝 문자열의 endOfWord를 false로 바꾸는 것 역시 올바른 delete 동작이다.

 

이와 달리,

문자열의 끝에 도달했으나 endOfWord = false인 경우

문자열의 부분 문자가 자식 노드로 존재하지 않는 경우

 

이는 곧 if(false)로 만들며, 최종적으로 false를 반환해 해당 문자열이 트라이에 저장되지 않아서 delete가 실패했음을 의미한다.