基于重复平方法的2阶矩阵n次幂Python代码优化（n>30万耗时超10s）

首先明确：重复平方（快速幂）的算法本身时间复杂度为O(logn)，哪怕n达到1e18规模，也仅需要约60次矩阵乘法计算，n=30万的场景理论运算耗时应在毫秒级，你遇到的超时问题和算法选型无关，完全是代码实现层面的问题。

原代码的核心问题 #

• 逻辑错误
  • 矩阵乘法中result[1][1]的计算漏了交叉项Y[1][0] * X[0][1]，输出结果本身不符合矩阵乘法的数学定义
  • n=0时返回整数1，而非矩阵幂对应的单位矩阵[[1,0],[0,1]]，后续运算会触发类型不匹配错误
  • 沿用Python2风格的/做整数除法，在Python3环境下会生成浮点数作为递归参数，当n超过2^53的浮点数精度阈值时，会因精度丢失导致递归终止条件失效，甚至出现无限递归
• 性能损耗点
  • 递归实现的快速幂有额外开销：Python的函数调用、栈帧开辟成本远高于普通循环，且递归逻辑没有复用已计算的平方结果，存在冗余乘法
  • 嵌套列表存储的寻址开销：每次访问矩阵元素需要做两层列表索引查找，每次计算乘法还要临时创建嵌套列表存结果，这部分操作在Python中的耗时占比极高
  • 大整数运算的隐式开销：如果计算过程中不对矩阵元素取模，元素的位数会随幂次呈指数增长，n=30万对应的矩阵元素可达数万位十进制长度，两个大整数相乘的耗时会随位数快速上升，这是你超时的最核心原因

针对性优化方案 #

• 先修正所有基础逻辑bug
  • 补全矩阵乘法的计算项，修正0次幂的单位矩阵返回值
  • 把除法替换为Python3兼容的整数地板除法//，避免隐式的浮点转换
• 递归快速幂改写为迭代实现
  • 消除函数调用的额外开销，按n的二进制位从低位到高位遍历，将结果初始化为单位矩阵，每轮遍历先将基矩阵平方，若当前二进制位为1则将基矩阵乘入结果，全程无冗余乘法计算
• 优化矩阵的存储结构
  • 2×2矩阵不需要嵌套列表这类复杂结构，直接把四个元素拆成独立的局部变量存储，彻底消除列表索引的寻址开销（Python中局部变量的访问速度是嵌套列表索引访问的3~5倍），返回值直接用元组存储，省去临时列表的创建成本
• 按需加入模运算
  • 如果计算场景是线性递推求值（比如斐波那契数列、线性递推项），完全可以在每次矩阵乘法后对四个元素做取模操作，将所有参与运算的数限制在模数的取值范围内，从根本上避免大整数运算的性能损耗。哪怕n达到1e18，只要参与运算的是固定范围的小整数，单次乘法的耗时都是纳秒级

注意：不要为了2×2矩阵的计算引入numpy。对这种极小尺寸的矩阵，numpy的函数调用、数组结构调度的开销远大于纯Python直接做算术运算，反而会降低运行速度。

优化后的参考实现 #

def mul2x2(x00, x01, x10, x11, y00, y01, y10, y11, mod=None):
    # 计算两个2x2矩阵的乘积，返回四个独立元素
    r00 = x00 * y00 + x01 * y10
    r01 = x00 * y01 + x01 * y11
    r10 = x10 * y00 + x11 * y10
    r11 = x10 * y01 + x11 * y11
    if mod is not None:
        r00 %= mod
        r01 %= mod
        r10 %= mod
        r11 %= mod
    return r00, r01, r10, r11

def power_mat(M, n, mod=None):
    if n == 0:
        # 返回单位矩阵
        return (1, 0, 0, 1)
    # 结果初始化为单位矩阵
    res00, res01, res10, res11 = 1, 0, 0, 1
    # 提取输入矩阵的四个元素
    base00, base01, base10, base11 = M[0][0], M[0][1], M[1][0], M[1][1]
    if mod is not None:
        base00 %= mod
        base01 %= mod
        base10 %= mod
        base11 %= mod
    while n > 0:
        if n & 1: # 当前二进制位为1时，将基矩阵乘入结果
            res00, res01, res10, res11 = mul2x2(
                res00, res01, res10, res11,
                base00, base01, base10, base11,
                mod
            )
        # 基矩阵自乘平方
        base00, base01, base10, base11 = mul2x2(
            base00, base01, base10, base11,
            base00, base01, base10, base11,
            mod
        )
        n >>= 1 # 右移一位，处理下一个二进制位
    return [
        [res00, res01],
        [res10, res11]
    ]

优化效果 #

按上述实现，哪怕n取1e6，只要传入合适的mod参数将元素限制在小整数范围，整体耗时也不会超过1毫秒，远低于你设定的10秒阈值。如果不需要精确计算超大整数结果，这部分大整数带来的性能损耗是完全可以避免的。

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