LeetCode/problems/count-vowels-permutation.rs

use std::collections::HashMap;

type CacheKey = (i32, char);
struct Cache {
    cache: HashMap<CacheKey, i32>,
}

impl Cache {
    const MODULO: i32 = 1000000007;

    fn new() -> Cache {
        let mut cache = HashMap::<CacheKey, i32>::new();

        for c in "aeiou".chars() {
            cache.insert((1, c), 1);
        }

        Cache { cache: cache }
    }

    fn is_precomputed(&self, n: i32) -> bool {
        self.cache.keys().any(|key| key.0 == n)
    }

    fn get(&self, n: i32) -> i32 {
        self.cache
            .keys()
            .filter(|key| key.0 == n)
            .fold(0, |acc, key| {
                // Since the answer may be too large, return it modulo 10^9 + 7.
                (acc + self.cache.get(key).unwrap_or(&0)) % Self::MODULO
            })
    }

    fn precompute(&mut self, n: i32) {
        if self.is_precomputed(n) {
            return;
        }

        if !self.is_precomputed(n - 1) {
            self.precompute(n - 1);
        }

        for (current, continuation) in vec![
            ('a', "e"),
            ('e', "ai"),
            ('i', "aeou"),
            ('o', "iu"),
            ('u', "a"),
        ] {
            self.cache.insert(
                (n, current),
                continuation.chars().fold(0, |acc, vowel| {
                    (acc + self.cache.get(&(n - 1, vowel)).unwrap()) % Self::MODULO
                }),
            );
        }
    }
}

struct Solution {}
impl Solution {
    pub fn count_vowel_permutation(n: i32) -> i32 {
        let mut cache = Cache::new();
        cache.precompute(n);

        cache.get(n)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_example_1() {
        // Input: n = 1
        // Output: 5
        // Explanation: All possible strings are: "a", "e", "i" , "o" and "u".
        assert_eq!(count_vowel_permutation(1), 5);
    }

    #[test]
    fn test_example_2() {
        // Input: n = 2
        // Output: 10
        // Explanation: All possible strings are: All possible strings are: "ae", "ea", "ei", "ia", "ie", "io", "iu", "oi", "ou" and "ua".
        assert_eq!(count_vowel_permutation(2), 10);
    }

    #[test]
    fn test_example_3() {
        // Input: n = 5
        // Output: 68
        assert_eq!(count_vowel_permutation(5), 68);
    }

    #[test]
    fn test_example_4() {
        // Input: n = 3
        // Output: 19
        assert_eq!(count_vowel_permutation(3), 19);
    }
}