TrieStringSet.java 2.67 KB
import java.util.*;

public class TrieStringSet {
    private static class StringTrieNode {
        public Map<Character, StringTrieNode> children;
        public boolean isPresent;

        public StringTrieNode() {
            this.children = new HashMap<>();
            this.isPresent = false;
        }
    }
    private StringTrieNode root;
    private int size;

    public TrieStringSet() {
        this.size = 0;
        this.root = new StringTrieNode();
    }

    public int size() {
        return this.size;
    }

    public boolean contains(String elem) {
        StringTrieNode curr = this.root;
        while (curr != null && !elem.isEmpty() && curr.children.containsKey(elem.charAt(0))) {
            curr = curr.children.get(elem.charAt(0));
            elem = elem.substring(1);
        }
        return curr != null && elem.isEmpty() && curr.isPresent;
    }

    public void add(String elem) {
        add(elem, this.root);
    }

    private void add(String elem, StringTrieNode curr) {
        if (elem.isEmpty()) {
            curr.isPresent = true;
        } else {
            if(!curr.children.containsKey(elem.charAt(0))) {
                curr.children.put(elem.charAt(0), new StringTrieNode());
            }
            add(elem.substring(1), curr.children.get(elem.charAt(0)));

        }
    }

    public int numNodes() {
        return numNodes(this.root);
    }

    private int numNodes(StringTrieNode curr) {
        // count this node
        int count = 1;

        // count the subtrees
        for (char c : curr.children.keySet()) {
            count += numNodes(curr.children.get(c));
        }
        return count;
    }

    public int numWords() {
        return numWords(this.root);
    }

    private int numWords(StringTrieNode curr) {
        int count = 0;
        if (curr.isPresent) {
            count++;
        }
        for (char c : curr.children.keySet()) {
            count += numNodes(curr.children.get(c));
        }
        return count;
    }

    // Collect all of the words in the set into a list
    public List<String> getAllWords() {
        List<String> result = new ArrayList<>();
        getAllWords(result, this.root, "");
        return result;
    }

    private void getAllWords(List<String> result, StringTrieNode curr, String acc) {
        if (curr == null) {
            return;
        }
        if (curr.isPresent) {
            result.add(acc);
        }
        for (char c : curr.children.keySet()) {
            acc += c; // "choose"
            getAllWords(result, curr.children.get(c),acc);
            acc = acc.substring(0, acc.length() - 1); // "unchoose"
            // getAllWords(result, curr.children.get(c), acc + c);
        }

    }

}