从盒子移除物品的方式数高效求解问题（n≤1e8）

我遇到一道运行时约束严格（运行时长<10s、无大内存占用）的算法题，目前思路受阻，编写的解法无法通过半数测试用例。

题目描述 #

盒子中装有若干物品，每次仅能移除1个或3个物品，求清空盒子共有多少种方式。由于答案数值可能非常大，返回结果需对10^9+7取模。

例如初始物品数n=7时，共有9种移除方式，具体枚举如下：

1.(1,1,1,1,1,1,1)
2.(1,1,1,1,3)
3.(1,1,1,3,1)
4.(1,1,3,1,1)
5.(1,3,1,1,1)
6.(3,1,1,1,1)
7.(1,3,3)
8.(3,1,3)
9.(3,3,1)

n=7时正确返回值为9。

函数要求 #

实现的函数接收参数n代表物品总数，返回整数表示清空盒子的方式数。

约束范围 #

1<=n<=10^8

测试样例 #

Input: 1
Output: 1
Explanation: 仅有一种移除1个物品的方式，结果取模后为1

Input:7
Output:9
Explanation: 共9种移除7个物品的方式

首先推导得到的递推关系是正确的：当n>3时，f(n) = f(n-1) + f(n-3)，初始边界值f(1)=1, f(2)=1, f(3)=2。但原有实现存在三个致命问题，无法适配大n场景：

• 递归实现存在深度限制，Python默认递归深度约为1000，n超过该值会直接抛出递归深度溢出错误
• 即使改写为迭代递推，O(n)的时间复杂度在n=1e8时，Python解释器执行简单循环也需要数十秒，远超10s的时间限制；同时记忆化字典存储n个结果的内存开销也很高
• 取模操作放在最后执行，n较大时中间整数会膨胀到极高位数，大幅提升运算开销；另外math.pow返回浮点数，大数值下会出现精度丢失问题。

线性递推关系可以通过矩阵快速幂将时间复杂度降到O(logn)，对于n=1e8的场景仅需要20余次矩阵乘法运算，运行时间在毫秒级，内存占用可忽略。

原理 #

递推式f(n) = f(n-1) + f(n-3)对应的3阶转移矩阵为：
$$
M = \begin{bmatrix}
1 & 0 & 1\
1 & 0 & 0\
0 & 1 & 0
\end{bmatrix}
$$
满足如下关系：
$$
\begin{bmatrix}f(n)\f(n-1)\f(n-2)\end{bmatrix} = M^{n-3} \times \begin{bmatrix}f(3)\f(2)\f(1)\end{bmatrix} = M^{n-3} \times \begin{bmatrix}2\1\1\end{bmatrix}
$$
通过快速幂算法计算矩阵的幂，所有乘法、加法运算中途对10^9+7取模即可。

实现代码 #

MOD = 10**9 + 7

def matrix_mult(a, b):
    # 3x3矩阵乘法
    res = [[0]*3 for _ in range(3)]
    for i in range(3):
        for k in range(3):
            if a[i][k] == 0:
                continue
            for j in range(3):
                res[i][j] = (res[i][j] + a[i][k] * b[k][j]) % MOD
    return res

def matrix_pow(mat, power):
    # 矩阵快速幂，初始化为单位矩阵
    res = [[1 if i == j else 0 for j in range(3)] for i in range(3)]
    while power > 0:
        if power % 2 == 1:
            res = matrix_mult(res, mat)
        mat = matrix_mult(mat, mat)
        power = power // 2
    return res

def numberOfWays(n):
    if n == 1 or n == 2:
        return 1
    if n == 3:
        return 2
    M = [
        [1,0,1],
        [1,0,0],
        [0,1,0]
    ]
    mat = matrix_pow(M, n-3)
    # 乘初始向量[f(3),f(2),f(1)] = [2,1,1]
    return (mat[0][0] * 2 + mat[0][1] * 1 + mat[0][2] * 1) % MOD

性能验证 #

• n=1时返回1，n=7时返回9，符合样例要求
• n=1e8时运算次数仅为log2(1e8)≈27次矩阵乘法，总运算量不到千次，运行时长远小于1s，内存占用仅为几个3x3矩阵的空间，完全满足约束要求。

内容的提问来源于stack exchange，提问作者Wboy