slow trie

main.rs

// testing
extern crate itertools;
use itertools::Zip;
// fin testing

use std::io;
use std::io::*;
use std::ops::Deref;

use std::collections::HashMap;

#[derive(Debug, Clone)]
struct NaiveTrie {
    children: HashMap<char, Box<NaiveTrie>>,
}

impl NaiveTrie {
    fn new() -> Self {
        let mut empty_string_char = HashMap::default();
        empty_string_char.insert('^', Box::new(NaiveTrie { children: HashMap::default() }));

        NaiveTrie { children: empty_string_char }
    }

    /*
    fn max_children_count_in_node(&self) -> usize {
        let children = if self.children.get(&'$').is_some() {
            1
        } else {
            0
        };
        let childrens_children: usize = self.children
            .iter()
            .map(|(_, child)| child.deref().max_children_count_in_node())
            .sum();
        children + childrens_children
    }
    */

    fn max_children_count_in_node(&self) -> usize {
        let total: &mut usize = &mut 0;
        self.max_children_count_in_node_recursive(total);
        *total
    }

    fn max_children_count_in_node_recursive(&self, total: &mut usize) {
        if self.children.get(&'$').is_some() {
            *total += 1;
        }

        for (_, child) in &self.children {
            child.max_children_count_in_node_recursive(total);
        }
    }

    fn insert_str(&mut self, slice: &str) {
        let char_vec: &Vec<char> = &["^", slice].join("").chars().collect();
        self.insert_recursive(char_vec, 0);
    }

    fn insert_recursive(&mut self, char_vec: &Vec<char>, depth: usize) {
        if depth + 1 == char_vec.len() {
            let mut new_child_hashmap = HashMap::default();
            new_child_hashmap.insert('$', Box::new(NaiveTrie { children: HashMap::default() }));
            let new_child = &mut NaiveTrie { children: new_child_hashmap };
            self.children.insert(char_vec[depth], Box::new(new_child.clone()));
            return;
        }

        match self.children.get_mut(&char_vec[depth]) {
            // we have to write return because the "fail" in the borrow checker!!!
            Some(ref mut children) => {
                // println!("encontrado!");
                return children.insert_recursive(char_vec, depth + 1);
            }
            None => {
                // lines after match should go here when the borrow checker gets better
            }
        }

        let mut new_child_hashmap = HashMap::default();
        new_child_hashmap.insert(char_vec[depth + 1],
                                 Box::new(NaiveTrie { children: HashMap::default() }));

        let new_child = &mut NaiveTrie { children: new_child_hashmap };

        new_child.insert_recursive(char_vec, depth + 1);
        self.children.insert(char_vec[depth], Box::new(new_child.clone()));
    }

    fn count_words_with_prefix(&self, prefix: &str) -> usize {
        let char_vec = ["^", prefix].join("").chars().collect();
        self.count_words_with_prefix_recursive(&char_vec, 0)
    }

    // necesito el máximo número de hijos del árbol en el punto del prefijo
    fn count_words_with_prefix_recursive(&self,
                                         char_vec: &Vec<char>,
                                         current_char_pos: usize)
                                         -> usize {
        return match self.children.get(&char_vec[current_char_pos]) {
                   Some(child) => {
                       if current_char_pos + 1 >= char_vec.len() {
                           child.as_ref().max_children_count_in_node()
                       } else {
                           child.count_words_with_prefix_recursive(char_vec, current_char_pos + 1)
                       }
                   }
                   None => 0,
               };
    }
}

fn main() {
    let mut naive_trie = NaiveTrie::new();

    /*
    // let words_to_insert = vec!["hack", "hackerrank",];
    //let words_to_insert = vec!["s", "ss", "sss", "ssss", "sssss"];
    let words_to_insert = vec!["pepe", "pepa", "pepito", "juanito"];

    for word in &words_to_insert {
        naive_trie.insert_str(word);
    }
    */

    /*
    for word in &words_to_insert[..1] {
        println!("{}", naive_trie.count_words_with_prefix(word));
    }
    */

    // println!("{:#?}", &naive_trie);

    let mut buffer = String::new();
    io::stdin().read_to_string(&mut buffer).expect("Failed to read input");

    let lines: Vec<&str> = buffer.split("\n").collect();

    for line in &lines[1..] {
        let line_parts: Vec<&str> = line.split(' ').collect();
        if line_parts[0] == "add" {
            naive_trie.insert_str(line_parts[1]);
        } else {
            println!("{}", naive_trie.count_words_with_prefix(line_parts[1]));
        }
    }
}

#[test]
fn test_trie1() {
    let mut naive_trie = NaiveTrie::new();
    let words_to_insert = vec!["hack", "hackerrank"];
    let prefixes = vec!["", "hac", "hak", "hackerrankazo"];
    let expected_outputs = [2, 2, 0, 0];

    for word in words_to_insert {
        naive_trie.insert_str(word);
    }

    for (prefix, expected_output) in Zip::new((&prefixes, &expected_outputs)) {
        assert_eq!(&naive_trie.count_words_with_prefix(prefix), expected_output);
    }
}

#[test]
fn test_trie2() {
    let mut naive_trie = NaiveTrie::new();
    let words_to_insert = vec!["pepe", "pepa", "pepito", "juanito"];
    let prefixes = vec!["", "pi", "pepi", "pep", "pe", "p", "juanit", "juanito", "j", "v"];
    let expected_outputs = vec![4, 0, 1, 3, 3, 3, 1, 1, 1, 0];

    for word in words_to_insert {
        naive_trie.insert_str(word);
    }

    for (prefix, expected_output) in Zip::new((&prefixes, &expected_outputs)) {
        assert_eq!(&naive_trie.count_words_with_prefix(prefix), expected_output);
    }
}

#[test]
fn test_trie3() {
    let mut naive_trie = NaiveTrie::new();
    let words_to_insert = vec!["s", "ss", "sss", "ssss", "sssss"];
    let prefixes = vec!["s", "ss", "sss", "ssss", "sssss", "ssssss"];
    let expected_outputs = vec![5, 4, 3, 2, 1, 0];

    for word in words_to_insert {
        naive_trie.insert_str(word);
    }

    for (prefix, expected_output) in Zip::new((&prefixes, &expected_outputs)) {
        assert_eq!(&naive_trie.count_words_with_prefix(prefix), expected_output);
    }
}