Cod sursa(job #2237420)

Utilizator a_h1926Heidelbacher Andrei a_h1926 Data 1 septembrie 2018 20:06:08
Problema Trie Scor 55
Compilator java Status done
Runda Arhiva educationala Marime 2.9 kb
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;

public final class Main {
  public static final String IN_FILE = "trie.in";
  public static final String OUT_FILE = "trie.out";

  public static final class Trie {
    public static final int SIGMA = 26;

    private static final class Node {
      final Node[] sons = new Node[SIGMA];
      int size = 0;
      int count = 0;
    }

    private static int encode(final char symbol) {
      return (int) symbol - (int) 'a';
    }

    private final Node root = new Node();

    public void insert(final String word) {
      Node node = root;
      for (int i = 0; i < word.length(); i++) {
        final int s = encode(word.charAt(i));
        node.size++;
        if (node.sons[s] == null) {
          node.sons[s] = new Node();
        }
        node = node.sons[s];
      }
      node.size++;
      node.count++;
    }

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

    private boolean erase(final Node node, final String word, final int i) {
      node.size--;
      if (i == word.length()) {
        node.count--;
      } else {
        final int s = encode(word.charAt(i));
        if (erase(node.sons[s], word, i + 1)) {
          node.sons[s] = null;
        }
      }
      return node.size == 0;
    }

    public int count(final String word) {
      Node node = root;
      for (int i = 0; i < word.length(); i++) {
        final int s = encode(word.charAt(i));
        if (node.sons[s] == null) {
          return 0;
        }
        node = node.sons[s];
      }
      return node.count;
    }

    public int lcp(final String word) {
      Node node = root;
      int prefix = 0;
      for (int i = 0; i < word.length(); i++) {
        final int s = encode(word.charAt(i));
        if (node.sons[s] == null) {
          return prefix;
        }
        prefix++;
        node = node.sons[s];
      }
      return prefix;
    }
  }

  private static BufferedReader newReader(final String fileName)
      throws IOException {
    return new BufferedReader(
        new InputStreamReader(new FileInputStream(fileName)));
  }

  public static void main(final String[] args) throws IOException {
    try (final BufferedReader reader = newReader(IN_FILE);
        final PrintWriter writer = new PrintWriter(OUT_FILE)) {
      final Trie trie = new Trie();
      String line = reader.readLine();
      while (line != null) {
        final String[] tokens = line.split(" ");
        final int type = Integer.parseInt(tokens[0]);
        final String word = tokens[1];
        if (type == 0) {
          trie.insert(word);
        } else if (type == 1) {
          trie.erase(word);
        } else if (type == 2) {
          writer.println(trie.count(word));
        } else {
          writer.println(trie.lcp(word));
        }
        line = reader.readLine();
      }
    }
  }
}