import java.util.*;

/*
To avoid duplicates with with skips we need to find start nodes that are as close as possible to the root
If it is not possible to find a node for a prefix with a given amount of errors, ONLY then we try to apply skips
 */

class Solution {

    static final int ERRORS_ALLOWED = 2;

    static class TrieNode {
        Map<Character, TrieNode> map = new HashMap<>();
        String word = null;
    }

    TrieNode root = new TrieNode();

    public Solution(List<String> dictionary) {
        for (String word : dictionary) {
            TrieNode cur = root;
            for (char c : word.toCharArray()) {
                if (!cur.map.containsKey(c)) {
                    TrieNode next = new TrieNode();
                    cur.map.put(c, next);
                }
                cur = cur.map.get(c);
            }
            cur.word = word;
        }
    }

    List<String> getWords(String prefix) {
        List<TrieNode> startNodes = getStartNodes(root, prefix, 0, 0, new HashSet<>());
        List<String> words = new ArrayList<>();
        for (TrieNode node : startNodes) {
            words.addAll(getAllWords(node));
        }
        return words;
    }

    private List<String> getAllWords(TrieNode node) {
        List<String> words = new ArrayList<>();
        if (node.word != null) {
            words.add(node.word);
        }
        for (TrieNode next : node.map.values()) {
            words.addAll(getAllWords(next));
        }
        return words;
    }

    private List<TrieNode> getStartNodes(TrieNode node, String prefix, int ind, int errors, Set<TrieNode> visited) {
        List<TrieNode> res = new ArrayList<>();
        if (visited.contains(node)) {
            return res;
        }
        if (ind >= prefix.length()) {
            visited.addAll(traverse(node));
            return Collections.singletonList(node);
        }
        char charOfPrefix = prefix.charAt(ind);
        if (node.map.containsKey(charOfPrefix)) {
            res.addAll(getStartNodes(node.map.get(charOfPrefix), prefix, ind + 1, errors, visited));
        }
        if (errors < ERRORS_ALLOWED) {
            for (char c : node.map.keySet()) {
                if (node.map.containsKey(charOfPrefix) && c == charOfPrefix) {
                    continue;
                }
                List<TrieNode> withPrefixShift = getStartNodes(node.map.get(c), prefix, ind + 1, errors + 1, visited);
                res.addAll(withPrefixShift);
                res.addAll(getStartNodes(node.map.get(c), prefix, ind, errors + 1, visited));
            }
        }
        return res;
    }

    private List<TrieNode> traverse(TrieNode node) {
        List<TrieNode> treeNodes = new ArrayList<>();
        treeNodes.add(node);

        for (TrieNode next : node.map.values()) {
            treeNodes.addAll(traverse(next));
        }

        return treeNodes;
    }

    public static void main(String[] args) {
        testSkipNotNeeded();
        testSkipNeeded();
    }


    static void testSkipNotNeeded() {
        List<String> dictionary = List.of("abc", "doecm", "abzd", "azfs", "kzcs");
        Solution solution = new Solution(dictionary);
        List<String> words = solution.getWords("abc");
        System.out.println(words);
    }
    
    static void testSkipNeeded() {
        List<String> dictionary = List.of("abzd", "helloaaa", "heoaabb", "heloaabb");
        Solution solution = new Solution(dictionary);
        List<String> words = solution.getWords("heoaa");
        System.out.println(words);
    }
}
