1. Introduction
You asked:
what should I do at first when I face this kind of problem?
This is a good question! When I studied computing at university I was taught the following general technique:
Analyze the program and work out its complexity — that is, how does its runtime depend on its input?
Measure its runtime on some small examples.
Combine the complexity from step (1) with the measurements from step (2) to estimate the runtime for the worst case.
If the worst case is unacceptable, try to find a better algorithm.
As you get better at carrying out these steps, some of them will become second nature to you and so you won't have to go through them all in such detail as I do below. But when you're starting out, or unsure about the approach to take it's good practice to be thorough.
2. Analysis
The way to analyze a program is to annotate each line with the time it takes and the number of times it is executed. Then all you have to do is multiply and add!
But that's easier said than done, because how on earth do we know how long each line takes to execute? If we really wanted an accurate answer we could estimate it by running lots of tests. But we don't actually need to compute an exact answer: we just want an idea of how the runtime varies as a function of the input size. So as a first approximation, we'll say that simple operations (assignment, calling xrange, etc) take constant time 1, and that arithmetic operations like addition and subtraction take time proportional to the number of digits in the numbers. (We'll see later that this is an over-simplification, but it will have to do for now.)
CODE TIME TAKEN NUMBER OF EXECUTIONS
============================ ========== ====================
sum = 0 1 1
z = k 1 1
for i in xrange(1, n+1, 2): 1 n / 2
d = z * 2 log z n / 2
sum += d log sum n / 2
z = z * k log z n / 2
if (n + 1) % 2 == 0: log n 1
sum = sum - d // 2 log sum 0 or 1
r = sum % 1000000007 log sum 1
(Note that I've rewritten the code to use n and k instead of a and b, so that the code more closely matches the statement of the problem.)
Now, in order to turn this into a total runtime, we need to know how big sum and z become. Well, \$z\$ gets multiplied by \$k\$ on each iteration of the loop, so after the \$i\$th iteration, \$z = k^{i+1}\$, and so \$z\$ is no more than about \$k^{n\over 2}\$ and so \$\log z\$ is no more than about \$n \log k \over 2\$.
At each iteration we add \$2z\$ to sum, where \$z\$ is no more than \$k^{n\over 2}\$. And there are about \$n \over 2\$ loop iterations, so sum is no more than \$nk^{n\over 2}\$ and so \$\log{sum}\$ is no more than \${n\log k\over2} + \log n\$. So adding everything up, we get a total runtime bounded by $$1 + 1 + {n\over 2} + {n^2 \log k \over 4} + {n^2 \log k + n\log n \over 2} + {n^2 \log k \over 4} + \log n + 2{n \log k + \log n \over 2}$$ which we can reorganize in terms of \$n\$ to get $$ n^2 \log k + {n\log n\over 2} + (\log k + {1\over 2})n + 2\log n + 2.$$ Now, as \$n\$ gets large, the terms containing \$n^2\$ will dominate (all the other terms become small by comparison). So we expect that as \$n\$ becomes large the runtime will become roughly proportional to \$n^2\log k\$.
(If this looks kind of complicated, it's because I've written out all the steps in detail. When you get experienced with the method, then you can skip all the steps that you know aren't going to appear in the result and just concentrate on the "big guns". For example, I would say to myself, "hmmm, that's \$Θ(n)\$ loops, each of which is handling numbers of typical size \$Θ(n\log k)\$ so the total runtime is \$Θ(n^2\log k)\$.")
Measurement is easier if we organize the code into functions, like this:
def palindromes(n, k):
"""Return the number of palindromes up to n letters long, using an
alphabet of k letters, modulo 1000000007.
"""
sum = 0
z = k
for i in xrange(1, n + 1, 2):
d = z * 2
sum += d
z = z * k
if (n + 1) % 2 == 0:
sum = sum - d // 2
return sum % 1000000007
Now we can measure the time taken (using Python's built-in timeit module) for some small values of \$n\$:
from timeit import timeit
def t(n, k):
"""Return the time taken to compute palindromes(n, k)."""
return timeit(lambda:palindromes(n, k), number=1)
>>> t(10**3, 26)
0.0003402099828235805
>>> t(10**4, 26)
0.01810335897607729
>>> t(10**5, 26)
1.3598145439755172
>>> t(10**6, 26)
137.12631450797198
Let's plot the measurements from step (2) and draw a line of best fit to the equation \$t=an^2\$ that we worked out in step (1):
The fit looks pretty good and the coefficient of proportionality is \$a = 1.37 × 10^{-10}\$.
The worst case from the problem description is \$n=10^9\$, so this would result in a runtime of \$1.37 × 10^8\$ seconds, which is more than 4 years.
Four years is likely to exceed the time limit, so we need a better algorithm. How do you come up with one? Well, there are basically two sensible approaches here.
The first is to look for existing algorithms that solve the problem (or similar enough problems). It helps here to be familiar with lots of different types of algorithm so that you have some idea of where to look. And the way to become familiar with lots of algorithms is to read algorithms textbooks and then do some of the exercises.
The second approach is to use mathematical analysis to find alternative ways of computing the same results. For example, if you had a function like this:
def nsum(n):
"""Return a sum that depends on n."""
result = 0
for i in xrange(n):
result += 1
return result
Then I'm sure you could find a mathematical expression for the result in terms of the input \$n\$, and that would allow you to simplify the body of this function to:
return max(n, 0)
or even:
return n
if you didn't care about negative values of \$n\$.
3. Improving the algorithm
In this problem the mathematical approach looks like the way to go. Take a look at the analysis in §2.1 above. You can see that I was able to get quick estimates for the size of z and sum. But with a bit more care, could we work out the exact values for these variables?
Let's try that. We know that at the start of the \$i\$th iteration, \$z = k^i\$, and so \$d = 2k^i\$. So after \$m\$ loop iterations, the sum is $$ \sum_{1 \le i \le m} 2k^i = 2\sum_{1 \le i \le m} k^i. $$ If \$k=1\$ then this is \$2m\$, otherwise using the formula for summation of a geometric series it's $$ 2\left({k^{m + 1} - 1 \over k - 1} - 1\right). $$ There are exactly \$m = \left\lfloor{n + 1 \over 2}\right\rfloor\$ loop iterations, and if \$n\$ is odd we have a final subtraction of \$- k^m \$. So this suggests the following code:
def palindromes2(n, k, p=1000000007):
"""Return the number of palindromes up to n letters long, using an
alphabet of k letters, modulo p.
"""
if k == 1:
return n % p
m = (n + 1) // 2
r = 2 * ((k**(m + 1) - 1) // (k - 1) - 1)
if n % 2:
r -= k**m
return r % p
But we're not done! The first thing to do is to check that this is right, as it would have been easy to make a mistake in the mathematics. So let's check this function against the original:
>>> all(palindromes(n, k) == palindromes2(n, k)
... for n in range(1, 1000) for k in range(1, 27))
True
But is it fast enough? Measurements suggest that it is faster that the first version of the code, but not fast enough to get the answer for \$n=10^9\$ in reasonable time:
>>> t(10**4, 26)
0.00031872402178123593
>>> t(10**5, 26)
0.011414641980081797
>>> t(10**6, 26)
0.2926187649718486
>>> t(10**7, 26)
11.10361772798933
We had better go through the analysis again for the new program and estimate the runtime. This time I'll omit the detailed analysis and cut straight to the point: we have some arithmetic operations here, on numbers that get as big as \$k^{{n + 1 \over 2}+1}\$, that is, with roughly \$n \log k\$ digits. We expect additions and subtractions on numbers of this size to take time roughly proportional to \$n \log k\$, and similarly for multiplication and division by small numbers. But there are also two exponentations: how long do these take? Well, I don't know about you, but I'm not sure of the runtime complexity of Python's ** operator. So there are two ways to go here: we could dive into the details of Python's exponentiation algorithm (yuck), or we could just assume that computing \$k^m\$ takes time proportional to some power of \$n \log k\$, say \$1 + b\$, and work out the value of \$b\$ by fitting the measurements to the equation \$t = an^{1+b}\$. Here's the best fit:
with parameters \$a = 4 × 10^{-11}\$ and \$b=0.63\$. (The value for \$1+b\$ is close to the exponent in the Karatsuba multiplication algorithm, which is the algorithm that CPython uses for multiplication of big numbers — see the function k_mul in longobject.c.)
So in the worst case (\$n = 10^9\$) the new algorithm is going to take about 5 hours. This is better than the 4 years of the original algorithm but is still likely to be unacceptable!
4. Further improving the algorithm
So we have to find another improvement. This is looking tricky: where can that come from? Well, the key observation is that the final result has to be returned modulo \$10^9+7\$, and so it can never have more than ten digits. And yet the intermediate values that we are computing are huge: in the worst case, \$k^m\$ has hundreds of millions of digits. This is a terrible waste of computation if we are going to compute a number with at most ten digits. So is there a way that we can do all the arithmetic modulo \$10^9+7\$ and so keep the intermediate values small?
How is that going to work? Modular addition, subtraction and multiplication are easy — we just add, subtract or multiply normally and then take the modulus later. For efficient modular exponentiation we can use the built-in function pow.
But we need a modular division operation in order to divide by \$k - 1\$, and Python doesn't have one built-in. In fact, modular division doesn't always exist — we can only divide by \$x\$ modulo \$p\$ if there is a unique number \$y\$ such that \$xy ≡ 1 \pmod p\$. (In this case \$y\$ is called the "multiplicative inverse of \$x\$ modulo \$p\$".) Luckily the modulus \$p\$ is prime, and this means that the multiplicative inverse \$y\$ exists for every \$x ≠ 0\$, and we can use the extended Euclidean algorithm to compute it.
That leads to this implementation:
def modular_inverse(x, p):
"""Return the multiplicative inverse of x modulo the prime p."""
y, z, a, b = 0, 1, p, x
while b:
a, (q, b) = b, divmod(a, b)
y, z = z, y - q * z
return y
def palindromes3(n, k, p=1000000007):
"""Return the number of palindromes up to n letters long, using an
alphabet of k letters, modulo the prime p.
"""
if k == 1:
return n % p
m = (n + 1) // 2
r = 2 * ((pow(k, m + 1, p) - 1) * modular_inverse(k - 1, p) - 1)
if n % 2:
r -= pow(k, m, p)
return r % p
Again, we must check that this is correct:
>>> all(palindromes(n, k) == palindromes3(n, k)
for n in range(1, 1000) for k in range(1, 27))
True
and see how bad the worst case is:
>>> timeit(lambda:palindromes3(10**9, 26), number=1)
3.3604970667511225e-05
That's probably acceptable now.
