Cod sursa(job #3265563)

Utilizator lucametehauDart Monkey lucametehau Data 31 decembrie 2024 14:29:25
Problema Trie Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 3.14 kb
use std::fs::File;
use std::io::*;

#[derive(Clone)]
struct Trie {
    end_count: i32,
    children_count: i32,
    children: Vec<Option<Box<Trie>>>,
}

impl Trie {
    fn new() -> Trie {
        Trie {
            end_count: 0,
            children_count: 0,
            children: vec![None; 26],
        }
    }

    fn insert(&mut self, word: &str, ind: usize) {
        if ind == word.len() { 
            self.end_count += 1;
            return;
        }
        let char_index = (word.as_bytes()[ind] - b'a') as usize;
        if self.children[char_index].is_none() {
            self.children_count += 1;
            self.children[char_index] = Some(Box::new(Trie::new()));
        }
        self.children[char_index].as_mut().unwrap().insert(word, ind + 1);
    }

    fn insert_word(&mut self, word: &str) {
        self.insert(word, 0);
    }

    fn erase(&mut self, word: &str, ind: usize) -> bool {
        if ind == word.len() {
            self.end_count -= 1;
        }
        else {
            let char_index = (word.as_bytes()[ind] - b'a') as usize;
            assert_eq!(self.children[char_index].is_none(), false);

            if self.children[char_index].as_mut().unwrap().erase(word, ind + 1) {
                self.children[char_index] = None;
                self.children_count -= 1;
            }
        }

        return self.end_count == 0 && self.children_count == 0;
    }

    fn erase_word(&mut self, word: &str) {
        let _ = self.erase(word, 0);
    }

    fn get_count(&self, word: &str, ind: usize) -> i32 {
        if ind == word.len() { return self.end_count; }
        
        let char_index = (word.as_bytes()[ind] - b'a') as usize;
        return if let Some(child) = &self.children[char_index] { child.get_count(word, ind + 1) } else { 0 };
    }

    fn count_appearances(&self, word: &str) -> i32 {
        self.get_count(word, 0)
    }

    fn get_lcp(&self, word: &str, ind: usize) -> usize {
        if ind == word.len() { return ind; }
        
        let char_index = (word.as_bytes()[ind] - b'a') as usize;
        return if let Some(child) = &self.children[char_index] { child.get_lcp(word, ind + 1) } else { ind };
    }

    fn lcp(&self, word: &str) -> usize {
        self.get_lcp(word, 0)
    }
}

fn main() {
    let in_file = File::open("trie.in").unwrap();
    let out_file = File::create("trie.out").unwrap();

    let input = BufReader::new(in_file);
    let mut output = BufWriter::new(out_file);

    let mut root = Trie::new();

    for line_result in input.lines() {
        let line = line_result.unwrap();
        // println!("{}", line);
        let operation: Vec<&str> = line.split_whitespace().collect();
        let (query_type, word) = (operation[0].parse::<i32>().unwrap(), operation[1]);

        // println!("operation is {}; word is {}", query_type, word);

        if query_type == 0 {
            root.insert_word(word);
        }
        else if query_type == 1 {
            root.erase_word(word);
        }
        else if query_type == 2 {
            writeln!(output, "{}", root.count_appearances(word)).unwrap();
        }
        else {
            writeln!(output, "{}", root.lcp(word)).unwrap();
        }
    }
}