Cod sursa(job #2999168)

Utilizator popoiu.georgeGeorge Popoiu popoiu.george Data 10 martie 2023 15:59:25
Problema Trie Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 2.7 kb
use std::{
    fs::File,
    error::Error,
    io::{BufReader, BufRead, LineWriter, Write},
};

#[derive(Debug, Clone)]
struct Trie {
    cnt: i32,
    fii: [Option<Box<Trie>>; 26],
    nrfii: i32,
}

impl Default for Trie {
    fn default() -> Trie {
        Trie {
            cnt: 0,
            fii: Default::default(),
            nrfii: 0,
        }
    }
}

impl Trie {
    fn add(&mut self, word: &str) {
        if word.len() == 0 {
            self.cnt += 1;
            return;
        }        
        
        let c = (word.as_bytes()[0] - 97) as usize;
        if let None = self.fii[c] {
            self.fii[c] = Some(Box::new(Trie::default()));
            self.nrfii += 1;
        }        
        let next = self.fii[c].as_mut().unwrap();                             
        next.add(&word[1..]);
    }
    fn del(&mut self, word: &str) {
        if word.len() == 0 {
            self.cnt -= 1;
            return;
        }
        let c = (word.as_bytes()[0] - 97) as usize;        
        let next = self.fii[c].as_mut().unwrap();        
        next.del(&word[1..]);
        if next.cnt == 0 && next.nrfii == 0 {
            self.nrfii -= 1;
            self.fii[c] = None;                
        }
    }
    fn count(&mut self, word: &str) -> i32 {
        if word.len() == 0 {
            return self.cnt;
        }
        let c = (word.as_bytes()[0] - 97) as usize;
        if let None = self.fii[c] {
            return 0;
        }
        let next = self.fii[c].as_mut().unwrap();
        return next.count(&word[1..]);
    }
    fn longest_prefix(&mut self, word: &str) -> i32 {
        if word.len() == 0 {
            return 0;
        }
        let c = (word.as_bytes()[0] - 97) as usize;
        if let Some(next) = self.fii[c].as_mut() {
            return 1 + next.longest_prefix(&word[1..]);
        }        
        0
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("trie.in")?;
    let out_file = File::create("trie.out")?;
    let mut writer = LineWriter::new(out_file);
    let reader = BufReader::new(file);    
    let mut trie = Trie::default();
    // println!("{:?}", t);
    for line in reader.lines() {
        // writeln!(writer, "{}", line?)?;
        let line_s = line?;  
        if line_s.len() == 0 {
            continue;
        }      
        let mut it = line_s.split_whitespace();
        let op: i32 = it.next().unwrap().parse()?;
        let word = it.next().unwrap();
        if op == 0 {
            trie.add(word);
        } else if op == 1 {
            trie.del(word);
        } else if op == 2 {
            writeln!(writer, "{}", trie.count(word))?;
        } else {
            writeln!(writer, "{}", trie.longest_prefix(word))?;
        }
    }
    Ok(())
}