Python序列求和问题：HackerRank代码部分测试未通过排查

问题背景 #

数学层面定义：序列第i项为 a_i = i² - (i-1)²，化简后可得 a_i = 2i - 1，即该序列是[1, 2n]范围内的所有奇数。需求是实现函数计算前n项和Sₙ，并返回Sₙ mod (10^9+7)的结果。

原代码问题 #

你编写的代码如下：

def summingSeries(n):
    # Firstly, an anonimous function computing the ith term in the sequence.
    a_i = lambda i: 2*i - 1
    
    # Now, let us sum all elements in the list consisting of
    # every a_i for i in the range [1, n].
    s_n = sum([a_i(x) for x in range(1, n + 1)])

    # Lastly, return the required modulo.
    return s_n % (10**9 + 7)

该代码仅通过部分测试用例，核心问题在于效率不足：

• 当n是极大值（比如10^18量级）时，range(1, n+1)无法生成这么长的序列，会直接触发内存溢出；
• 即使能生成序列，sum遍历整个列表的时间复杂度是O(n)，对于超大n来说，计算时间会远超题目限制，导致超时。

优化方案（数学推导） #

前n个奇数的和有现成的数学结论：Sₙ = n²。比如：

• n=1时，1=1²
• n=2时，1+3=4=2²
• n=3时，1+3+5=9=3²
  直接利用这个结论，就能把时间复杂度降到O(1)，完全适配所有测试用例。

优化后的代码 #

def summingSeries(n):
    MOD = 10**9 + 7
    return (n * n) % MOD

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