1220. 统计元音字母序列的数目 :「线性 DP」&「矩阵快速幂」

1,211 阅读6分钟

题目描述

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

Tag : 「线性 DP」、「矩阵快速幂」

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

字符串中的每个字符都应当是小写元音字母('a', 'e', 'i', 'o', 'u')

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

由于答案可能会很大,所以请你返回 模 109+710^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:

输入:n = 5

输出:68

提示:

  • 1<=n<=21041 <= n <= 2 * 10^4

线性 DP

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

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

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

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

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

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

代码(利用对称性统计方案数代码见 P2P2,感谢 @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)
};
  • 时间复杂度:令 CC 为字符集大小,本题 CC 固定为 55。整体复杂度为 O(nC)O(n * C)
  • 空间复杂度:O(nC)O(n * C)

矩阵快速幂

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

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

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

  1. 简单题学「矩阵快速幂」
  2. 简单题学「矩阵快速幂」Ⅱ
  3. 矩阵快速幂的运用 :从「状态机 DP」到「矩阵快速幂」

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

ans=[f[n1][0]f[n1][1]f[n1][2]f[n1][3]f[n1][4]]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=[f[0][0]f[0][1]f[0][2]f[0][3]f[0][4]]=[11111]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][X] 依赖于 f[i1][X]f[i - 1][X],同时在「解法一」中明确了各个 f[i][j]f[i][j]f[i1][X]f[i - 1][X] 的关系。

根据「矩阵乘法」,不难发现:

[f[i][0]f[i][1]f[i][2]f[i][3]f[i][4]]=[0100010100110110010110000][f[i1][0]f[i1][1]f[i1][2]f[i1][3]f[i1][4]]\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=[0100010100110110010110000]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=matmat...matorians = mat * mat * ... * mat * ori

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

ans=matn1orians = 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
}
  • 时间复杂度:共需要执行 logn\log{n} 次矩阵操作,运算量最大的矩阵操作为 matmatmat * mat,复杂度为 C3C^3。整体复杂度为 O(lognC3)O(\log{n} * C^3)
  • 空间复杂度:复杂度上界为 matmat 大小,复杂度为 O(C2)O(C^2)

最后

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

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

为了方便各位同学能够电脑上进行调试和提交代码,我建立了相关的仓库:github.com/SharingSour…

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