Cod sursa(job #2916542)

Utilizator CiobaCatalinCioba Catalin CiobaCatalin Data 30 iulie 2022 14:20:25
Problema Trie Scor 0
Compilator java Status done
Runda Arhiva educationala Marime 2.85 kb
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;

public class Main {
    public static void main(String[] args) throws IOException {
        File in = new File("trie.in");
        File out = new File("trie.out");

        Scanner sc = new Scanner(in);
        PrintWriter pw = new PrintWriter(out);

        TrieNode trie = new TrieNode();

        while (sc.hasNext()) {
            int type = sc.nextInt();
            String word = sc.next();

            if (type == 0) {
                trie.add(word);
            } else if (type == 1) {
                trie.delete(word);
            } else if (type == 2) {
                pw.println(trie.count(word));
            } else {
                pw.println(trie.prefix(word));
            }
        }

        pw.close();
        sc.close();
    }

    private static class TrieNode {
        int endingHere;
        int count;
        private Map<Character, TrieNode> children;

        public TrieNode() {
            endingHere = 0;
            count = 0;
            children = new HashMap<>();
        }

        public void add(String word) {
            TrieNode current = this;

            for (char c : word.toCharArray()) {
                TrieNode node = current.children.getOrDefault(c, new TrieNode());
                current.children.put(c, node);
                current.count++;
                current = node;
            }

            current.count++;
            current.endingHere++;
            // System.out.println(word + " -> " + current.endingHere);
        }

        public void delete(String word) {
            TrieNode current = this;

            for (char c : word.toCharArray()) {
                current.count--;
                current = current.children.get(c);
            }

            current.count--;
            current.endingHere--;
            // System.out.println(word + " -> " + current.endingHere);
        }

        public int count(String word) {
            TrieNode current = this;

            for (char c : word.toCharArray()) {
                current = current.children.get(c);
            }

            // System.out.println(word + " found " + current.endingHere);
            return current.endingHere;
        }

        public int prefix(String word) {
            TrieNode current = this;
            int prefix = 0;

            for (char c : word.toCharArray()) {
                TrieNode node = current.children.get(c);
                if (node != null && node.count > 0) {
                    prefix++;
                    current = node;
                } else {
                    // System.out.println("Stopping at " + c);
                    break;
                }
            }
            // System.out.println(word + " prefix " + prefix);
            return prefix;
        }
    }
}