题目描述
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.
Example 1:
Input: n = 1
Output: 5
Explanation: All possible strings are: “a”, “e”, “i” , “o” and “u”.
Example 2:
Input: n = 2
Output: 10
Explanation: All possible strings are: “ae”, “ea”, “ei”, “ia”, “ie”, “io”, “iu”, “oi”, “ou” and “ua”.
Example 3:
Input: n = 5
Output: 68
Constraints:
- 1 <= n <= 2 * 10^4
题目分析
由于规则中对后续字符有特定要求，也就是说，如果首字母为a，那么长度为2的字符串仅有ae这一种；即相同首字母的字符串，长度L+1和L之间存在关系，由题目中的规则约束。那么，这个问题我们尝试使用dp解决
我们用dp[i][j]表示以字符i开头，且长度为j的字符串数量，那么，根据规则，则有:
dp[a][j] = dp[e][j-1] // ‘a’ may only be followed by ‘e’
dp[e][j] = dp[a][j-1] + dp[i][j-1] // ‘e’ may only be followed by ‘a’ or ‘i’
… // 其他以规则类推
按照以上思路，代码如下
// AC, need optimize
func countVowelPermutation(n int) int {
mod := 1000000007
ans := 0
vowels := []byte{'a', 'e', 'i', 'o', 'u'}
// 规则map
rules := map[byte][]byte{
'a': {'e'},
'e': {'a', 'i'},
'i': {'a', 'e', 'o', 'u'},
'o': {'i', 'u'},
'u': {'a'},
}
// m即dp，表示以不同字母开头的不同长度的字符串个数
m := make(map[byte]map[int]int, 5)
for _, v := range vowels {
m[v] = make(map[int]int)
// base case
m[v][1] = 1
}
for i := 2; i <= n; i++ {
for _, v := range vowels {
count := 0
// 按照规则求解
for _, iv := range rules[v] {
count += m[iv][i-1] % mod
}
m[v][i] = count % mod
}
}
// 最后，将长度为n的结果取模加和即可
for _, v := range m {
ans += v[n] % mod
}
return ans % mod
}
以上代码提交AC后，查看题目Hint，发现提示的dp定义与刚才的思路不同；这里，题目提示的是，dp[i][j]表示长度为j以i结尾的字符串数量。按照这个思路，需要对规则做一点转换；我们需要由原始规则推导出： 如果当前字符为x，能够出现在x之前的字符集合
按照这个思路，解法如下:
// dp 16ms
// Let dp[i][j] be the number of strings of length i that ends with the j-th vowel.
func countVowelPermutationV2(n int) int {
mod := 1000000007
dp := make([][]int, n+1)
for i := 0; i <= n; i++ {
dp[i] = make([]int, 5)
}
// base case
dp[1] = []int{1, 1, 1, 1, 1}
// 按照原规则推导出的能够出现在当前字符之前的字符集合
// 0~4 分别表示 a e i o u
pres := map[int][]int{
0: {1, 2, 4},
1: {0, 2},
2: {1, 3},
3: {2},
4: {2, 3},
}
// dp[i][j] = sum(dp[i-1][x]), x∈[vowels which can appear before j]
for i := 2; i <= n; i++ {
for j := 0; j < 5; j++ {
count := 0
for _, v := range pres[j] {
count += dp[i-1][v] % mod
}
dp[i][j] = count % mod
}
}
ans := 0
for i := 0; i < 5; i++ {
ans += dp[n][i] % mod
}
return ans % mod
}
上述代码继续提交AC后，发现还不是最快的。看了一下更快的代码，思路还是第一种解法的dp思路，只是不再维护一个很大的dp数组，而是使用一个一维数组，然后在迭代计算过程中不断更新 具体解法如下:
优化代码
func countVowelPermutationV3(n int) int {
mod := 1000000007
// base case
// 这里使用了1维数组，然后在计算过程中不断更新此数组；这样就避免了很多中间结果占用内存
dp := []int{1, 1, 1, 1, 1}
for i := 2; i <= n; i++ {
tmp := make([]int, 5)
tmp[0] += dp[1] % mod
tmp[1] += dp[0]%mod + dp[2]%mod
tmp[2] += dp[0]%mod + dp[1]%mod + dp[3]%mod + dp[4]%mod
tmp[3] += dp[2]%mod + dp[4]%mod
tmp[4] += dp[0] % mod
dp = tmp
}
ans := 0
for _, v := range dp {
ans += v % mod
}
return ans % mod
}