幾何級数

初項1、公比r、項数nの等比級数1+r+r^2+...+r^(n-1)を効率的に求めたい。
何も工夫せずにやると

def f(r,n):
  ans = 0
  b = 1
  for i in xrange(n):
    ans += b
    b *= r
  return ans

でO(n)となる。
よく知られている公式を用いると次のようにも書ける。

def f(r,n):
  return (1 - pow(r,n)) / (1 - r)

これだとべき乗の計算にバイナリ法を用いることでO(log n)となるが、実用上はmodをとって考えることが多く、modと1-rが互いに素でないときに逆元を求めることができず困ってしまう(逆に言えば、modが素数だったりで逆元を計算可能ならこれで十分。ただしr=1の場合に注意)。それ以外にも行列だと逆行列を求める必要があったりで困る。蟻本ではこの問題を行列累乗で解いていて、それでまったく不便はないのだけど今回は別の方法を考えてみる。
n=2*kのときを考える。

• 1+r+r^2+...+r^(2*k-1)
• = (1+r^2+...+r^(2*k-2)) + (r+r^3+...+r^(2*k-1))
• = (1+r^2+...+r^(2*k-2)) + r*(1+r^2+...+r^(2*k-2))
• = (1 + r) * (1+r^2+...+r^(2*k-2))
• = (1 + r) * (1+R+R^2+...+R^(k-1)) (R = r^2)

と再帰的に計算できる。nが奇数のときは1+r*sum(r,n-1)となるので問題ない。コードに落とすとこうなる。

def f(r,n):
  if n == 0 : return 0
  if n%2 == 1 : return 1 + r*f(r,n-1)
  return (1+r) * f(r*r, n/2)

個人的にはこっちの方がバイナリ法っぽくて好き。というかそのものなので計算量はO(log n)。多分剰余計算が少ないので実用上も行列累乗より早いことが期待できる。