我应该如何优化这段代码？

Question

def f(n):
    Total_Triangles = 0
    for i in range(1,n+1):
        term = 3**(i-1)
        Total_Triangles+=term
    return Total_Triangles
Q = int(input())
for i in range(Q):
    n = int(input())
    Ans = f(n)*4 +1
    print(Ans%1000000007)

如何解决此代码中的时间限制错误？

Answer 1

卡兰有一个很好的答案。它会加快你原来的方法，但你最终还是会计算出巨大的数字。幸运的是，Python 的 Long 类型可以做到这一点，但我认为它不如原生 32 位或 64 位整数类型高效。

你被告知要给答案模一个巨大的数字 M，1,000,000,007。你可以通过使用modular arithmetic 来改进算法，这样你的数字永远不会变得很大。在模算术中，这是真的：

(a + b) % M == (a % M + b % M) % M
(a * b) % M == (a % M * b % M) % M

一种方法是使用模运算预先计算所有可能的 Q 值：

M = 1000000007

def makef(m):
    """Generator to create all sum(3**i) mod M"""

    n = 1
    s = 0

    for i in range(m):
        yield s

        s = (s + n) % M
        n = ((n  + n) % M + n) % M

f = list(makef(100000))
Q = int(input())

for i in range(Q):
    n = int(input())
    print (f[n] * 4 + 1) % M

这将在一个大循环中进行计算，但只有一次，并且应该足够快以满足您的要求。

Python 为您提供了第二种方式：表达式a ** b 映射到内置函数pow(a, b)。这个函数可以接受第三个参数：模运算的基数，这样pow(a, b, M) 将计算(a ** b) % M 而不会产生巨大的中间结果。

现在您可以使用 Karan 的简洁公式了。但是等等，有一个陷阱：你必须将幂的结果除以二。上述模数关系不适用于除法。例如，(12 // 2) % M 是 6，但如果你先应用模运算符，就像 pow 函数所做的那样，你会得到 ((12 % 2) // 2) % M，它是 1，而不是你想要的。一个解决方案是计算幂模2 * M，然后除以2：

def f(n):
    return pow(3, n, 2 * 1000000007) // 2

Q = int(input())
for i in range(Q):
    n = int(input())
    print (f(n) * 4 + 1) % M

（请注意，3 的所有幂都是奇数，所以我删除了 - 1 并让整数除法来完成工作。）

旁注：选择M 的值是为了使两个小于M 的数字相加适合带符号的32 位整数。这意味着 C、C++ 或 Java 的用户不必使用 bignum 库。但请注意，3 * n 仍然可以溢出有符号整数，因此在乘以 3 时必须小心：请改用 ((n + n) % M + n) % M。

Answer 2

你要找3 ** 0 + 3 ** 1 ... + 3 ** (n - 1)，这只是一个几何级数，第一项a = 1，公比r = 3，项数n = n，使用几何级数公式求和，我们可以找到f(n) 如此定义时要快得多：

def f(n):
    return (3 ** n - 1) // 2