1220. Count Vowels Permutation

Question

Given an integer n, your task is to count how many strings of length n can be formed under the following rules:

• Each character is a lower case vowel ('a', 'e', 'i', 'o', 'u')

• Each vowel 'a' may only be followed by an 'e'.

• Each vowel 'e' may only be followed by an 'a' or an 'i'.

• Each vowel 'i' may not be followed by another 'i'.

• Each vowel 'o' may only be followed by an 'i' or a 'u'.

• Each vowel 'u' may only be followed by an 'a'.

Since the answer may be too large, return it modulo 10^9 + 7.

Solution

DFS搜索，记忆化剪枝。
每次根据上一个字母last来从哈希表里选择下一个可选字母进行遍历，并计算总的可组成数sum。

Code

class Solution {
    final int MOD = (int) Math.pow(10, 9) + 7;
    int[][] memo;
    HashMap<Character, Character[]> map;
    public int countVowelPermutation(int n) {
        if(n == 1) return 5;
        map = new HashMap<>();
        memo = new int[26][n+1];
        Character[] vowels = {'a', 'e', 'i', 'o', 'u'}, 
        a = {'e'}, e = {'a','i'}, i = {'a','e','o','u'}, o = {'i', 'u'}, u = {'a'};
        map.put('a', a);
        map.put('e', e);
        map.put('i', i);
        map.put('o', o);
        map.put('u', u);
        
        int ans = 0;
        
        for(char vowel : vowels){
            ans += count(vowel, n-1);
            ans %= MOD;
        }
        return ans;
    }
    
    private int count(char last, int n){
        if(memo[last-'a'][n] != 0) return memo[last-'a'][n];
        if(n == 1){
            memo[last-'a'][n] = map.get(last).length;
            return memo[last-'a'][n];
        }
        
        int sum = 0;
        for(char next : map.get(last)){
            sum += count(next, n-1);
            sum %= MOD;
        }
        memo[last-'a'][n] = sum;
        return sum;
    }
}