Cod sursa(job #2998980)

Utilizator popoiu.georgeGeorge Popoiu popoiu.george Data 10 martie 2023 12:34:07
Problema Trie Scor 0
Compilator rs Status done
Runda Arhiva educationala Marime 2.6 kb
use std::{
    fs::File,
    error::Error,
    io::{BufReader, BufRead, LineWriter, Write},    
};

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

impl Default for Trie {
    fn default() -> Trie {
        Trie {
            cnt: 0,
            fii: Vec::<Trie>::new(),
            flag: [false; 26],
            nrfii: 0,
        }
    }
}

impl Trie {
    fn add(&mut self, word: &str) {
        if word.len() == 0 {
            self.cnt += 1;
            return;
        }
        if self.fii.len() == 0 {
            self.fii.resize(26, Trie::default());
        }
        let c = (word.as_bytes()[0] - 97) as usize;                
        if self.flag[c] == false {
            self.flag[c] = true;
            self.nrfii += 1;
        }
        let next = &mut self.fii[c];
        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 = &mut self.fii[c];
        next.del(&word[1..]);
        if next.cnt == 0 && next.nrfii == 0{
            self.flag[c] = false;
            self.nrfii -= 1;
        }
    }
    fn count(&mut self, word: &str) -> i32 {
        if word.len() == 0 {
            return self.cnt;
        }
        let c = (word.as_bytes()[0] - 97) as usize;
        let next = &mut self.fii[c];
        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 self.flag[c] == true {
            let next = &mut self.fii[c];
            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?;        
        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(())
}