大输入规模下Python幂次求和取模代码优化求助

计算规则 #

输入第一行包含三个空格分隔的整数N、x、n，取值范围为：1 < N ≤ 1000000、0 < x ≤ 10000、1 < n ≤ 10000。参数说明：

• N：最大幂次
• x：指定底数
• n：输出结果的模运算除数
  后续跟随N-1行输入，每行是一个不超过10000的非负整数。

计算示例

输入数据：

5 2 100
2
8
1
3

正确输出为84，计算逻辑如下：

2 * (2**2) = 8
8 * (2**3) = 64
1 * (2**4) = 16
3 * (2**5) = 96
8 + 64 + 16 + 96 = 184
184 % 100 = 84

待优化问题

现有Python实现在N取值较大时运行效率极低，代码如下：

N, x, n = map(int, input().split())
sum = 0
for power in range(2,N+1):
   sum+= (int(input())* (x**power))
   sum = sum%n 
print(sum)

优化方案

性能瓶颈分析

原有代码的性能问题主要来自两点：

1. 逐行调用input()读取输入，对于百万级的输入行数，IO开销极高。
2. 每次循环重新计算x**power，幂次升高后会产生超大整数运算，同时存在大量重复计算，时间复杂度达到O(N*log(power))。

优化思路

1. 输入批量读取：调用sys.stdin.read()一次性读取所有输入内容，再批量转换为整数处理，大幅降低IO开销。
2. 幂次递推计算：利用模运算同余性质，递推得到每一轮的x^power % n结果，避免重复计算和大数运算，时间复杂度降至O(N)。

优化后代码

import sys

def main():
    data = list(map(int, sys.stdin.read().split()))
    ptr = 0
    N, x, n = data[ptr], data[ptr+1], data[ptr+2]
    ptr += 3
    res = 0
    x_mod = x % n
    # 初始为x^2 mod n，对应第一个系数的幂次
    current_pow = (x_mod * x_mod) % n
    for _ in range(N-1):
        res = (res + data[ptr] * current_pow) % n
        current_pow = (current_pow * x_mod) % n
        ptr += 1
    print(res)

if __name__ == "__main__":
    main()

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