Cod sursa(job #2230302)

Utilizator radustn92Radu Stancu radustn92 Data 9 august 2018 17:39:20
Problema Trie Scor 70
Compilator java Status done
Runda Arhiva educationala Marime 6 kb
import java.io.*;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {

    public static void main(String[] args) throws IOException {
        InputStream inputStream = new FileInputStream("trie.in");

        try (InputReader inputReader = new InputReader(inputStream);
            BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter("trie.out"))) {
            TrieGraph graph = new TrieGraph();

            while (!inputReader.hasReachedEnd()) {
                int opType = inputReader.nextInt();
                String word = inputReader.nextToken();

                switch (opType) {
                    case 0:
                        graph.addWord(word);
                        break;
                    case 1:
                        graph.removeWord(word);
                        break;
                    case 2:
                        bufferedWriter.write(graph.getWordCount(word) + "\n");
                        break;
                    case 3:
                        bufferedWriter.write(graph.getLongestCommonPrefix(word) + "\n");
                        break;
                    default:
                        throw new RuntimeException("Encountered unexpected opType");
                }
            }
        }
    }

    static class TrieGraph {

        public static final int SIGMA = 26;

        private TrieNode root;

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

        public void addWord(String word) {
            TrieNode currNode = root;
            currNode.updatePassingNodesCount(1);
            for (char entry : word.toCharArray()) {
                currNode = currNode.getOrCreateChild(entry - 'a');
                currNode.updatePassingNodesCount(1);
            }

            currNode.updateWordsEndingHere(1);
        }

        public void removeWord(String word) {
            removeWord(word.toCharArray(), 0, root);
        }

        public int getWordCount(String word) {
            TrieNode currNode = root;
            for (char entry : word.toCharArray()) {
                currNode = currNode.getChild(entry - 'a');
                if (currNode == null) {
                    return 0;
                }
            }

            return currNode.getWordsEndingHere();
        }

        public int getLongestCommonPrefix(String word) {
            TrieNode currNode = root;
            int commonPrefix = 0;
            for (char entry : word.toCharArray()) {
                currNode = currNode.getChild(entry - 'a');
                if (currNode == null) {
                    break;
                }
                commonPrefix++;
            }

            return commonPrefix;
        }

        public void removeWord(char[] word, int currPos, TrieNode currNode) {
            currNode.updatePassingNodesCount(-1);
            if (currPos == word.length) {
                currNode.updateWordsEndingHere(-1);
                return;
            }
            int childIdx = word[currPos] - 'a';
            TrieNode child = currNode.getChild(childIdx);

            removeWord(word, currPos + 1, child);
            if (child.getWordsPassingThroughThisNode() == 0) {
                currNode.removeChild(childIdx);
            }
        }
    }

    static class TrieNode {

        private TrieNode[] children;

        private int wordsEndingHere;

        private int wordsPassingThroughThisNode;

        public TrieNode() {
            wordsEndingHere = wordsPassingThroughThisNode = 0;
            children = new TrieNode[TrieGraph.SIGMA];
            Arrays.fill(children, null);
        }

        public void updatePassingNodesCount(int count) {
            wordsPassingThroughThisNode += count;
        }

        public void updateWordsEndingHere(int count) {
            wordsEndingHere += count;
        }

        public TrieNode getOrCreateChild(int childIdx) {
            TrieNode child = children[childIdx];
            if (child == null) {
                child = new TrieNode();
                children[childIdx] = child;
            }

            return child;
        }

        public int getWordsPassingThroughThisNode() {
            return wordsPassingThroughThisNode;
        }

        public int getWordsEndingHere() {
            return wordsEndingHere;
        }

        public TrieNode getChild(int childIdx) {
            return children[childIdx];
        }

        public void removeChild(int childIdx) {
            children[childIdx] = null;
        }
    }

    static class InputReader implements AutoCloseable {

        private BufferedReader bufferedReader;

        private StringTokenizer stringTokenizer;

        public InputReader(InputStream inputStream) {
            bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
            stringTokenizer = null;
        }

        public boolean hasReachedEnd() {
            return !canPrepareNextToken();
        }

        private boolean canPrepareNextToken() {
            if (stringTokenizer != null && stringTokenizer.hasMoreTokens()) {
                return true;
            }
            String nextLine;
            try {
                nextLine = bufferedReader.readLine();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            if (nextLine != null) {
                stringTokenizer = new StringTokenizer(nextLine);
            }

            return nextLine != null;
        }

        public String nextToken() {
            if (stringTokenizer == null || !stringTokenizer.hasMoreTokens()) {
                try {
                    stringTokenizer = new StringTokenizer(bufferedReader.readLine());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            return stringTokenizer.nextToken();
        }

        public int nextInt() {
            return Integer.parseInt(nextToken());
        }

        @Override
        public void close() throws IOException {
            bufferedReader.close();
        }
    }
}