LC 1220. 统计元音字母序列的数目

题目描述

这是 LeetCode 上的 1220. 统计元音字母序列的数目 ，难度为困难。

给你一个整数 n，请你帮忙统计一下我们可以按下述规则形成多少个长度为 n 的字符串：

字符串中的每个字符都应当是小写元音字母（’a’, ‘e’, ‘i’, ‘o’, ‘u’）

每个元音 'a' 后面都只能跟着 'e'
每个元音 'e' 后面只能跟着 'a' 或者是 'i'
每个元音 'i' 后面不能再跟着另一个 'i'
每个元音 'o' 后面只能跟着 'i' 或者是 'u'
每个元音 'u' 后面只能跟着 'a'

由于答案可能会很大，所以请你返回模 $10^9 + 7$ 之后的结果。

示例 1：

输入：n = 1

输出：5

解释：所有可能的字符串分别是："a", "e", "i" , "o" 和 "u"。

示例 2：

输入：n = 2

输出：10

解释：所有可能的字符串分别是："ae", "ea", "ei", "ia", "ie", "io", "iu", "oi", "ou" 和 "ua"。

示例 3：

1
2
3

输入：n = 5

输出：68

提示：

$1 <= n <= 2 * 10^4$

线性 DP

定义 $f[i][j]$ 为考虑长度为 $i + 1$ 的字符串，且结尾元素为 $j$ 的方案数（其中 $j$ 代表数组 ['a', 'e', 'i', 'o', 'u'] 下标）。

不失一般性考虑 $f[i][j]$ 该如何计算。

我们可以从题意给定的规则进行出发，从 $f[i]$ 出发往前更新 $f[i + 1]$，也可以直接利用对称性进行反向分析。

为了方便大家理解，还是将常规的「从 $f[i]$ 出发往前更新 $f[i + 1]$」作为主要分析方法吧。

根据条件可以容易写出转移方程：

每个元音 'a' 后面都只能跟着 'e' ：$f[i + 1][1] += f[i][0]$；
每个元音 'e' 后面只能跟着 'a' 或者是 'i' ：$f[i + 1][0] += f[i][1]$、$f[i + 1][2] += f[i][1]$；
每个元音 'i' 后面不能再跟着另一个 'i' ：$f[i + 1][j] += f[i][2], (j 不能为 2)$；
每个元音 'o' 后面只能跟着 'i' 或者是 'u' ：$f[i + 1][2] += f[i][3]$、$f[i + 1][4] += f[i][3]$；
每个元音 'u' 后面只能跟着 'a' ：$f[i + 1][0] += f[i][4]$。

代码（利用对称性统计方案数代码见 $P2$，感谢 @Benhao 和 @5cm/s 🌸 提供的、我根本看不懂的、试图把困难题也做成「全鱼宴」题解的代码 🤣 ）：

class Solution {
    int MOD = (int)1e9+7;
    public int countVowelPermutation(int n) {
        long[][] f = new long[n][5];
        Arrays.fill(f[0], 1);
        for (int i = 0; i < n - 1; i++) {
            // 每个元音 'a' 后面都只能跟着 'e'
            f[i + 1][1] += f[i][0]; 
            // 每个元音 'e' 后面只能跟着 'a' 或者是 'i'
            f[i + 1][0] += f[i][1];
            f[i + 1][2] += f[i][1];
            // 每个元音 'i' 后面 不能 再跟着另一个 'i'
            f[i + 1][0] += f[i][2];
            f[i + 1][1] += f[i][2];
            f[i + 1][3] += f[i][2];
            f[i + 1][4] += f[i][2];
            // 每个元音 'o' 后面只能跟着 'i' 或者是 'u'
            f[i + 1][2] += f[i][3];
            f[i + 1][4] += f[i][3];
            // 每个元音 'u' 后面只能跟着 'a'
            f[i + 1][0] += f[i][4];
            for (int j = 0; j < 5; j++) f[i + 1][j] %= MOD;
        }
        long ans = 0;
        for (int i = 0; i < 5; i++) ans += f[n - 1][i];
        return (int)(ans % MOD);
    }
}

class Solution {
    int MOD = (int)1e9+7;
    public int countVowelPermutation(int n) {
        long[][] f = new long[n][5];
        Arrays.fill(f[0], 1);
        for (int i = 1; i < n; i++) {
            f[i][0] = f[i - 1][1];
            f[i][1] = f[i - 1][0] + f[i - 1][2];
            f[i][2] = f[i - 1][0] + f[i - 1][1] + f[i - 1][3] + f[i - 1][4];
            f[i][3] = f[i - 1][2] + f[i - 1][4];
            f[i][4] = f[i - 1][0];
            for (int j = 0; j < 5; j++) f[i][j] %= MOD;
        }
        long ans = 0;
        for (int i = 0; i < 5; i++) ans += f[n - 1][i];
        return (int)(ans % MOD);
    }
}

MOD = int(1e9) + 7
class Solution:
    def countVowelPermutation(self, n: int) -> int:
        f = [[1] * 5] + [[0] * 5 for _ in range(n - 1)]
        for i in range(1, n):
            f[i][0] = f[i-1][1]
            f[i][1] = (f[i-1][0] + f[i-1][2]) % MOD
            f[i][2] = ((((f[i-1][0] + f[i-1][1]) % MOD + f[i-1][3]) % MOD) + f[i-1][4]) % MOD
            f[i][3] = (f[i-1][2] + f[i-1][4]) % MOD
            f[i][4] = f[i-1][0]
        return (((((f[-1][0] + f[-1][1]) % MOD + f[-1][2]) % MOD + f[-1][3]) % MOD) + f[-1][4]) % MOD

const MOD int = 1e9 + 7
func countVowelPermutation(n int) int {
    f := make([][]int, n)
    f[0] = []int{1, 1, 1, 1, 1}
    for i := 1; i < n; i++ {
        f[i] = make([]int, 5) 
        f[i][0] = f[i-1][1]
        f[i][1] = (f[i-1][0] + f[i-1][2]) % MOD
        f[i][2] = (((f[i-1][0] + f[i-1][1]) % MOD + f[i-1][3]) % MOD + f[i-1][4]) % MOD
        f[i][3] = (f[i-1][2] + f[i-1][4]) % MOD
        f[i][4] = f[i-1][0]
    }
    return ((((f[n-1][0] + f[n-1][1]) % MOD + f[n-1][2]) % MOD + f[n-1][3]) % MOD + f[n-1][4]) % MOD
}

defmodule Solution do
  @spec count_vowel_permutation(n :: integer) :: integer
  def count_vowel_permutation(n) do
    solve(n, 1, 1, 1, 1, 1)
  end
  def solve(1, a, e, i, o, u) do
    rem((a + e + i + o + u), 1000000007)
  end
  def solve(n, a, e, i, o, u) do
    solve(n - 1, add(e, i, u), add(a, i), add(e, o), i, add(i, o))
  end
  def add(n1, n2) do
    add(n1, n2, 0)
  end
  def add(n1, n2, n3) do
    rem(n1 + n2 + n3, 1000000007)
  end
end

-spec count_vowel_permutation(N :: integer()) -> integer().
count_vowel_permutation(N) ->
  solve(N, 1, 1, 1, 1, 1).

solve(1, A, E, I, O, U) ->
  (A + E + I + O + U) rem 1000000007;
solve(N, A, E, I, O, U) ->
  solve(N - 1, add(E, I, U), add(A, I), add(E, O), I, add(I, O)).
add(N1, N2) ->
  add(N1, N2, 0).
add(N1, N2, N3) ->
  (N1 + N2 + N3) rem 1000000007.

(define/contract (count-vowel-permutation n)
  (-> exact-integer? exact-integer?)
    (define MOD 1000000007)
    (define (ad nums [tot 0])
      (cond
        [(empty? nums) (remainder tot MOD)]
        [else (ad (cdr nums) (+ tot (car nums)))]
      )
    )
    (let loop([n n] [a 1] [e 1] [i 1] [o 1] [u 1])
      (cond
        [(= n 1) (ad (list a e i o u))]
        [else (loop (- n 1) (ad (list e i u)) (ad (list a i)) (ad (list e o)) i (ad (list i o)))]
      )
    )
    )

const MOD = 10 ** 9 + 7;
const add = (...nums) => nums.reduce((a, b) => a + b) % MOD
var countVowelPermutation = function(n,a=1,e=1,i=1,o=1,u=1) {
    return n === 1 
        ? add(a, e, i, o, u)
        : countVowelPermutation(n - 1, add(e, i, u), add(a, i), add(e, o), i, add(i, o))
};

const MOD = 10 ** 9 + 7;
const add = (...nums) => nums.reduce((a, b) => a + b) % MOD
var countVowelPermutation = function(n) {
    let [a, e, i, o, u] = [1, 1, 1, 1, 1]
    while(--n)
        ([a, e, i, o, u] = [add(e, i, u), add(a + i), add(e + o), i, add(i, o)])
    return add(a, e, i, o, u)
};

时间复杂度：令 $C$ 为字符集大小，本题 $C$ 固定为 $5$。整体复杂度为 $O(n * C)$
空间复杂度：$O(n * C)$

矩阵快速幂

通常考虑递推能否使用「矩阵快速幂」来加速，主要依据为该递推过程是否为存在「结合律」线性递推过程。

如果是存在「结合律」的线性数列递推，诸如「斐波那契数列」等，都能使用「矩阵快速幂」来加速数列递推。

对「矩阵快速幂」不了解的同学，可以先看「矩阵快速幂の三部曲」（学过三部曲的同学做这道题应该是降维打击）：

我们最终需要求的是 $\sum_{i = 0}^{4} f[n - 1][i]$，将需要求得的部分列成列向量：

$ans = \begin{bmatrix} f[n - 1][0]\\ f[n - 1][1]\\ f[n - 1][2]\\ f[n - 1][3]\\ f[n - 1][4] \end{bmatrix}$

同时我们有起始的矩阵（每个元音字母独立成为一个字符）：

$ori = \begin{bmatrix} f[0][0]\\ f[0][1]\\ f[0][2]\\ f[0][3]\\ f[0][4] \end{bmatrix} = \begin{bmatrix} 1\\ 1\\ 1\\ 1\\ 1 \end{bmatrix}$

我们知道 $f[i][X]$ 依赖于 $f[i - 1][X]$，同时在「解法一」中明确了各个 $f[i][j]$ 与 $f[i - 1][X]$ 的关系。

根据「矩阵乘法」，不难发现：

$\begin{bmatrix} f[i][0]\\ f[i][1]\\ f[i][2]\\ f[i][3]\\ f[i][4] \end{bmatrix} = \begin{bmatrix} 0& 1& 0& 0& 0\\ 1& 0& 1& 0& 0\\ 1& 1& 0& 1& 1\\ 0& 0& 1& 0& 1\\ 1& 0& 0& 0& 0 \end{bmatrix} * \begin{bmatrix} f[i - 1][0]\\ f[i - 1][1]\\ f[i - 1][2]\\ f[i - 1][3]\\ f[i - 1][4] \end{bmatrix}$

我们令：

$mat = \begin{bmatrix} 0& 1& 0& 0& 0\\ 1& 0& 1& 0& 0\\ 1& 1& 0& 1& 1\\ 0& 0& 1& 0& 1\\ 1& 0& 0& 0& 0 \end{bmatrix}$

根据递推关系，可得：

$ans = mat * mat * ... * mat * ori$

再根据矩阵乘法具有「结合律」，最终可得：

$ans = mat^{n - 1} * ori$

代码（感谢 @Benhao 同学提供的「行向量」写法）：

class Solution {
    int MOD = (int)1e9+7;
    long[][] mul(long[][] a, long[][] b) {
        int r = a.length, c = b[0].length, z = b.length;
        long[][] ans = new long[r][c];
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                for (int k = 0; k < z; k++) {
                    ans[i][j] += a[i][k] * b[k][j];
                    ans[i][j] %= MOD;
                }
            }
        }
        return ans;
    }
    public int countVowelPermutation(int n) {
        long[][] ans = new long[][]{
            {1}, {1}, {1}, {1}, {1}
        };
        long[][] mat = new long[][]{
            {0, 1, 0, 0, 0},
            {1, 0, 1, 0, 0},
            {1, 1, 0, 1, 1},
            {0, 0, 1, 0, 1},
            {1, 0, 0, 0, 0},
        };
        int x = n - 1;
        while (x != 0) {
            if ((x & 1) != 0) ans = mul(mat, ans);
            mat = mul(mat, mat);
            x >>= 1;
        }
        long sum = 0;
        for (int i = 0; i < 5; i++) sum += ans[i][0];
        return (int)(sum % MOD);
    }
}

import numpy as np
MOD = 10 ** 9 + 7
dtype = np.dtype('uint64')
class Solution:
    def countVowelPermutation(self, n: int) -> int:
        ans, mat = np.ones(5, dtype=dtype), np.array([[0, 1, 1, 0, 1], [1, 0, 1, 0, 0], [0, 1, 0, 1, 0], [0, 0, 1, 0, 0], [0, 0, 1, 1, 0]],dtype=dtype)
        n -= 1
        while n:
            if n & 1:
                ans = ans @ mat % MOD
            mat = mat @ mat % MOD
            n >>= 1
        return int(ans.sum()) % MOD

const MOD int = 1e9 + 7

func countVowelPermutation(n int) (res int) {
    mul := func(X, Y [][]int) [][]int{
        res := make([][]int, len(X))
        for i := 0;i < len(res);i++{res[i] = make([]int, len(Y[0]))}
        for i := 0;i < len(res);i++{
            for j := 0;j < len(res[0]);j++{
                for k := 0;k < len(Y);k++{
                    res[i][j] = (res[i][j] + X[i][k] * Y[k][j] % MOD) % MOD
                }
            }
        }
        return res
    }

    ans, mat := [][]int{{1, 1, 1, 1, 1}}, [][]int{{0, 1, 1, 0, 1}, {1, 0, 1, 0, 0}, {0, 1, 0, 1, 0}, {0, 0, 1, 0, 0}, {0, 0, 1, 1, 0}}
    n--
    for n > 0 {
        if(n & 1 != 0){
            tmp := mul(ans, mat)
            ans = tmp
        }
        tmpMat := mul(mat, mat)
        mat = tmpMat
        n >>= 1
    }
    for _, row := range ans {
        for _, cell := range row{
            res = (res + cell) % MOD
        }
    }
    return
}

时间复杂度：共需要执行 $\log{n}$ 次矩阵操作，运算量最大的矩阵操作为 $mat mat$，复杂度为 $C^3$。整体复杂度为 $O(\log{n} C^3)$
空间复杂度：复杂度上界为 $mat$ 大小，复杂度为 $O(C^2)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.1220 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。

线性 DP 矩阵快速幂

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议，转载请注明出处！

LC 539. 最小时间差上一篇

LC 382. 链表随机节点下一篇