《小R的'chi'子序列挑战》、《小U的数列因子挑战》 | 豆包MarsCode AI刷题

一、小R的'chi'子序列挑战

问题理解

我们需要计算长度为 n 的、仅由 'c', 'h', 'i' 三种字符组成的字符串中，有多少是可爱串。可爱串的定义是：

1. 包含子序列 "chi"。
2. 不包含子串 "chi"。

数据结构选择

我们可以使用动态规划（DP）来解决这个问题。定义一个三维数组 dp 来表示状态：

• dp[pos][state][last_char_idx] 表示前 pos 个字符中，已匹配到状态 state 的 "chi" 子序列，前两个字符组合为 last_char_idx 的方案数。

算法步骤

1. 初始化状态：

   • 动态规划数组 dp，初始状态为 dp[0][0][last_char_map[(-1, -1)]] = 1。

2. 状态转移：

   • 遍历每一个位置 pos，对于每一种 state 和 last_char_idx，尝试添加字符 'c'、'h'、'i'。
   • 避免形成子串 "chi"：如果前两个字符是 'c'、'h'，当前字符是 'i'，则会形成子串 "chi"，跳过这种情况。
   • 更新状态：根据当前字符和已匹配的状态，更新 state 和 last_char_idx。

3. 计算结果：

   • 最终结果为所有在位置 n，状态为 3（已匹配到子序列 "chi"）的方案数之和。

代码实现

def solution(n: int) -> int:
    MOD = 10**9 + 7
    # 枚举前两个字符的组合
    last_char_map = {}
    last_chars = []
    idx = 0
    for lc1 in range(-1, 3):
        for lc2 in range(-1, 3):
            last_chars.append((lc1, lc2))
            last_char_map[(lc1, lc2)] = idx
            idx += 1
    # 初始化 dp 数组
    dp = [[[0 for _ in range(16)] for _ in range(4)] for _ in range(n+1)]
    dp[0][0][last_char_map[(-1, -1)]] = 1  # 初始状态

    for pos in range(n):
        for state in range(4):
            for last_char_idx in range(16):
                lc1, lc2 = last_chars[last_char_idx]
                count = dp[pos][state][last_char_idx]
                if count == 0:
                    continue
                for c in range(3):  # 遍历字符 'c', 'h', 'i'
                    # 判断是否会形成子串 "chi"
                    if lc1 == 0 and lc2 == 1 and c == 2:
                        continue  # 跳过会形成 "chi" 的情况
                    new_lc = (lc2, c)
                    new_last_char_idx = last_char_map[new_lc]
                    # 更新匹配状态
                    if state == 0 and c == 0:
                        new_state = 1
                    elif state == 1 and c == 1:
                        new_state = 2
                    elif state == 2 and c == 2:
                        new_state = 3
                    else:
                        new_state = state
                    dp[pos+1][new_state][new_last_char_idx] = (dp[pos+1][new_state][new_last_char_idx] + count) % MOD

    # 计算结果
    result = 0
    for last_char_idx in range(16):
        result = (result + dp[n][3][last_char_idx]) % MOD
    return result

时间和空间复杂度

时间复杂度

• 外层循环遍历字符串长度 n，时间复杂度为 O(n)。
• 内层循环遍历状态 state 和 last_char_idx，时间复杂度为 O(4 * 16)。
• 最内层循环遍历字符 'c', 'h', 'i'，时间复杂度为 O(3)。
• 因此，总的时间复杂度为 O(n * 4 * 16 * 3) = O(192n)，即 O(n)。

空间复杂度

• 动态规划数组 dp 的大小为 O(n * 4 * 16)，即 O(64n)，即 O(n)。 好的，下面将按照问题理解、数据结构选择、算法步骤、代码实现、时间和空间复杂度的结构详细解答该问题。

问题理解

问题描述：

小U得到一个数列，数列的定义如下：

• 初始条件：
  • f(1) = a
  • f(2) = b
• 递推关系：
  • 对于第 i 项，数列满足：f(i) = f(i-1) * f(i-2) * c^d

任务：

计算出数列的第 n 项 f(n) 的因子数量。由于答案可能很大，请对 10^9 + 7 取模。

数据结构选择

为了高效地计算数列 f(n) 的因子数量，我们需要以下步骤：

1. 质因数分解：
   • 质因数分解 a、b 和 c，得到它们的质因数及对应的指数。
2. 指数累加：
   • 基于递推关系 f(i) = f(i-1) * f(i-2) * c^d，我们可以通过累加质因数的指数来得到 f(n) 的质因数。
3. 因子数量计算：
   • 一个数的因子数量等于其质因数指数加一的乘积。例如，16 = 2^4，所以因子数量为 (4 + 1) = 5。

为了实现上述步骤，我们可以使用以下数据结构：

• 字典 (dict) 或 Counter：
  • 来存储每个质因数及其对应的指数。

算法步骤

步骤1：质因数分解

• 定义一个函数 factor(x)，用于对数 x 进行质因数分解，返回一个字典 {质因数: 指数}。

步骤2：初始化数列

• 对 f(1) = a 和 f(2) = b 进行质因数分解，获取它们的质因数字典。
• 将质因数字典存储在一个列表 factors 中，factors[i] 表示 f(i+1) 的质因数字典。

步骤3：递推计算质因数

• 对于每个 i 从 3 到 n：
  • 计算 f(i) 的质因数字典：
    • f(i) 的质因数包括 f(i-1) 和 f(i-2) 的质因数。
    • 还需要加入 c^d 的质因数，即将 c 的质因数指数乘以 d 再加到 f(i) 的质因数中。
  • 将 f(i) 的质因数字典添加到 factors 列表中。

步骤4：计算因子数量

• 获取 f(n) 的质因数字典。
• 计算因子数量，即对所有质因数的指数加一后相乘，再对 10^9 + 7 取模。

代码实现

def solution(a: int, b: int, c: int, d: int, n: int) -> int:
    from collections import defaultdict
    import math

    MOD = 10**9 + 7

    def factor(x):
        """对x进行质因数分解，返回质因数及其指数的字典"""
        factors = defaultdict(int)
        if x < 1:
            return factors  # 1及以下没有质因数
        # 处理2的因数
        while x % 2 == 0:
            factors[2] += 1
            x //= 2
        # 处理奇数因数
        for i in range(3, int(math.isqrt(x)) + 1, 2):
            while x % i == 0:
                factors[i] += 1
                x //= i
        # 如果x本身是一个大质数
        if x > 1:
            factors[x] += 1
        return factors

    # 因数分解c一次，因为c^d在递推中会多次使用
    factors_c = factor(c)

    # 初始化数列factors，factors[i]表示f(i+1)的质因数字典
    factors = []

    # f(1) = a
    factors_f1 = factor(a)
    factors.append(factors_f1)

    # f(2) = b
    factors_f2 = factor(b)
    factors.append(factors_f2)

    # 递推计算f(3)到f(n)
    for i in range(2, n):
        factors_prev1 = factors[i - 1]  # f(i)
        factors_prev2 = factors[i - 2]  # f(i-1)
        # 新的f(i+1)的质因数字典
        new_factors = defaultdict(int)
        # 合并f(i)的质因数
        for prime, exp in factors_prev1.items():
            new_factors[prime] += exp
        # 合并f(i-1)的质因数
        for prime, exp in factors_prev2.items():
            new_factors[prime] += exp
        # 合并c^d的质因数
        for prime, exp in factors_c.items():
            new_factors[prime] += d * exp
        factors.append(new_factors)

    # 获取f(n)的质因数字典
    if n == 1:
        factors_fn = factors[0]
    else:
        factors_fn = factors[n - 1]

    # 计算因子数量
    num_factors = 1
    for prime, exp in factors_fn.items():
        num_factors = (num_factors * (exp + 1)) % MOD

    return num_factors

时间和空间复杂度

时间复杂度

• 质因数分解：
  • 对于每个数 a、b 和 c 进行质因数分解的时间复杂度为 O(√x)，其中 x 是待分解的数。
• 递推计算：
  • 对于每个 f(i)，需要合并 f(i-1) 和 f(i-2) 的质因数字典，并加入 c^d 的质因数，时间复杂度为 O(k)，其中 k 是质因数的数量（通常较小，视输入数而定）。
  • 总的递推时间复杂度为 O(n * k)。
• 因子数量计算：
  • 遍历所有质因数，时间复杂度为 O(k)。

总体时间复杂度：O(n * k + √a + √b + √c)，由于 k 通常较小，可以认为接近 O(n)。

空间复杂度

• 质因数存储：
  • 对于每个 f(i)，存储其质因数字典，空间复杂度为 O(n * k)。
• 其他变量：
  • 常数级别。

总体空间复杂度：O(n * k)，其中 k 是质因数的数量。

测试用例验证

让我们验证上述代码对给定的测试用例是否正确：

1. 测试用例1：

   print(solution(1, 2, 3, 4, 3) == 10)
   

   解释：

   • f(3) = f(2) * f(1) * 3^4 = 2 * 1 * 81 = 162
   • 162 = 2 * 3^4
   • 因子数量计算：(1 + 1) * (4 + 1) = 2 * 5 = 10

   输出：True（正确）

2. 测试用例2：

   print(solution(2, 3, 5, 2, 4) == 30)
   

   解释：

   • f(3) = f(2) * f(1) * 5^2 = 3 * 2 * 25 = 150
   • f(4) = f(3) * f(2) * 5^2 = 150 * 3 * 25 = 11250
   • 11250 = 2 * 3^2 * 5^4
   • 因子数量计算：(1 + 1) * (2 + 1) * (4 + 1) = 2 * 3 * 5 = 30

   输出：True（正确）

3. 测试用例3：

   print(solution(1, 1, 2, 1, 5) == 5)
   

   解释：

   • f(1) = 1
   • f(2) = 1
   • f(3) = f(2) * f(1) * 2^1 = 1 * 1 * 2 = 2
   • f(4) = f(3) * f(2) * 2^1 = 2 * 1 * 2 = 4
   • f(5) = f(4) * f(3) * 2^1 = 4 * 2 * 2 = 16
   • 16 = 2^4
   • 因子数量计算：4 + 1 = 5

   输出：True（正确）