본문 바로가기
자료구조

[자료구조] Trie 구현 with Java

by Heesu.lee 2020. 10. 12.

TrieNode

public class TrieNode {

    private final Map<Character, TrieNode> child;
    private boolean isLast;

    public TrieNode() {
        child = new HashMap<>();
        isLast = false;
    }

    public TrieNode insert(Character character) {
        return child.computeIfAbsent(character, c -> new TrieNode());
    }

    public boolean isLast() {
        return isLast;
    }

    public void setLast(boolean isLast) {
        this.isLast = isLast;
    }

    public boolean contains(Character character) {
        return child.containsKey(character);
    }

    public TrieNode get(Character character) {
        return child.get(character);
    }

    public boolean isEmpty() {
        return child.isEmpty();
    }

    public void remove(char character) {
        child.remove(character);
    }
}

 

Trie

public class Trie {

    private final TrieNode root;

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

    public void insert(String word) {
        TrieNode ptr = root;

        for (int i = 0; i < word.length(); ++i) {
            ptr = ptr.insert(word.charAt(i));
        }
        ptr.setLast(true);
    }

    public boolean contains(String word) {
        TrieNode prev = root;

        for (int i = 0; i < word.length(); ++i) {
            TrieNode next = prev.get(word.charAt(i));

            if (next == null) {
                return false;
            }
            prev = next;
        }
        return prev.isLast();
    }

    public void delete(String word) {
        delete(root, word, 0);
    }

    private void delete(TrieNode parent, String word, int idx) {
        char character = word.charAt(idx);

        if (!parent.contains(character)) {
            throw new Error("There is no [ " + word + " ] in this Trie");
        }

        TrieNode child = parent.get(character);
        idx ++;

        if (idx == word.length()) {
            if (!child.isLast()) {
                throw new Error("There is no [ " + word + " ] in this Trie");
            }
            child.setLast(false);

            if (child.isEmpty()) {
                parent.remove(character);
            }
        }
        else {
            delete(child, word, idx);

            if (!child.isLast() && child.isEmpty()) {
                parent.remove(character);
            }
        }
    }
}

댓글