Python快速幂易超时原因及算法竞赛代码优化咨询

问题背景 #

在算法竞赛解题时，因C++的long long无法满足大数需求且不能使用__int128，改用Python实现。算法时间复杂度符合要求，但始终触发TLE（超时），排查发现无论是自定义的快速幂函数qpow，还是Python内置的pow()，性能都无法达标。

对比代码如下：

Python原代码（超时，运行超2000ms） #

from math import *

def phi(a):
    ans = 1
    up = int(sqrt(a))
    for i in range(2, up + 1):
        if a % i == 0:
            a //= i
            ans *= (i - 1)

            while a % i == 0:
                a //= i
                ans *= i
    if a > 1:
        ans *= (a - 1)
    return ans

def qpow(a, b, m):
    ans = 1
    tmp = a
    while b:
        if b & 1:
            ans = ans * tmp % m;
        tmp = tmp * tmp % m
        b >>= 1
    return ans

def exphi(a, b, m):
    p = phi(m)

    while b > p and b != b % p + p:
        b = b %p + p
    return qpow(a, b, m)

b, n, c = map(int, input().split(" "))
re = exphi(b, n - 1, c)
re = (b - 1) * re % c

if  re == 0:
    re = c
print(re)

C++对比代码（通过，仅63ms） #

#include <iostream>
#include <cmath>

using namespace std;
#define int long long
int phi(int a)
{
    int ans = 1;
    for (int i = 2; i * i <= a; i ++)
    {
        if (a % i == 0)
        {
            a /= i;
            ans *= (i - 1);
            while (a % i == 0)  a /= i, ans *= i;
        }
    }

    if (a > 1)
        ans *= (a - 1);
    return ans;
}

int qpow(int a, int b, int m)
{
    int ans = 1;
    int tmp = a;
    while (b)
    {
        if (b & 1)
            ans = ans * tmp % m;
        tmp = tmp * tmp % m;
        b >>= 1;
    }
    return ans;
}

// a^b = a ^ (b % phi(m) + phi(m)) (mod m)

signed main()
{
    string b, n;
    long long nc;
    cin >> b >> n >> nc;

    int nb = 0, nn = 0;

    for (auto i : b)
        nb = (nb * 10 + (i - '0')) % nc;

    int p = phi(nc);
    int flag = 0;
    for (auto i : n)
    {
            nn = (nn * 10 + (i - '0'));
            if (nn >= p)
                flag = true;
            nn %= p;
    }
 
    if (flag) nn += p;
    int re = qpow(nb, nn - 1, nc);

    re = (nb - 1 + nc) % nc * re % nc;
    if (!re)
        re = nc;
    cout << re << '\n';
    return 0;
}

超时原因分析 #

1. 解释型语言的性能瓶颈：Python是解释型语言，循环、位运算、模运算等基础操作的执行效率远低于编译为机器码的C++，快速幂的循环逻辑会产生大量解释开销。
2. 大数转换的隐性开销：原Python代码直接将输入的超大数字转为int类型，内存占用和转换时间远高于C++直接按字符串处理取模的方式。
3. 自定义函数的额外开销：自定义的qpow函数每一步操作都需要Python解释器逐行执行，即使是内置pow(a,b,m)（底层C实现），面对超大指数时，Python的大数运算开销也远高于C++。

Python代码优化方法 #

1. 优化输入处理，避免大数转换 #

参考C++的思路，直接对输入字符串进行取模运算，跳过超大数转int的步骤：

def str_mod(s, mod):
    res = 0
    for c in s:
        res = (res * 10 + int(c)) % mod
    return res

def str_mod_with_phi(s, phi_val):
    res = 0
    flag = False
    for c in s:
        res = res * 10 + int(c)
        if res >= phi_val:
            flag = True
            res %= phi_val
    if flag:
        res += phi_val
    return res

2. 替换自定义快速幂为内置pow(a,b,mod) #

Python内置的pow(a,b,mod)由C实现，执行效率远高于自定义的Python循环版本，直接替换即可。

3. 优化欧拉函数实现 #

避免浮点运算，改用i*i <= a作为循环条件，减少不必要的计算：

def phi(a):
    ans = 1
    i = 2
    while i * i <= a:
        if a % i == 0:
            a //= i
            ans *= (i - 1)
            while a % i == 0:
                a //= i
                ans *= i
        i += 1
    if a > 1:
        ans *= (a - 1)
    return ans

4. 简化指数处理逻辑 #

原exphi中的循环判断冗余，直接判断指数是否大于欧拉值即可：

def exphi(a, b, m):
    p = phi(m)
    if b > p:
        b = b % p + p
    return pow(a, b, m)

5. 通用优化技巧 #

• 使用sys.stdin.readline替代input()，提升输入速度；
• 尽量使用局部变量，Python访问局部变量的速度比全局变量更快；
• 小逻辑内联，减少函数调用的额外开销。

优化后的完整Python代码 #

import sys

def phi(a):
    ans = 1
    i = 2
    while i * i <= a:
        if a % i == 0:
            a //= i
            ans *= (i - 1)
            while a % i == 0:
                a //= i
                ans *= i
        i += 1
    if a > 1:
        ans *= (a - 1)
    return ans

def str_mod(s, mod):
    res = 0
    for c in s:
        res = (res * 10 + int(c)) % mod
    return res

def str_mod_with_phi(s, phi_val):
    res = 0
    flag = False
    for c in s:
        res = res * 10 + int(c)
        if res >= phi_val:
            flag = True
            res %= phi_val
    if flag:
        res += phi_val
    return res

def main():
    b_str, n_str, c_str = sys.stdin.readline().split()
    c = int(c_str)
    b = str_mod(b_str, c)
    
    if n_str == "1":
        re = (b - 1) % c
    else:
        p = phi(c)
        nn = str_mod_with_phi(n_str, p)
        nn -= 1
        if nn < 0:
            nn += p
        re = pow(b, nn, c)
        re = (b - 1) * re % c
    
    if re == 0:
        re = c
    print(re)

if __name__ == "__main__":
    main()

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