数列の推測

Project Eulerなどの問題を解いていると、「mod による場合分けを含んだ多項式により定義された数列」なんだろうけど、一般項を求めるのは面倒だなって思うことが良くある気がします。そこで、そのようなタイプの数列の一般項をすばやく求めるためのプログラムをpythonで書いてみました。
使い方は簡単です。a_nを求めたい場合はf(a,n)がa_nを返してくれます。

class Est:
  def __init__(self):
    self.mz = None
    pass

  def build(self,x):
    n = len(x)
    for k in xrange(1,n):
      z=[x[i::k] for i in xrange(k)]
      mz =[ self.f(x2) for x2 in z ]
      if all( len(mz[i])+2 <= len(z[i])  for i in xrange(k)):
        break
    else:
      print "Failed!"
    self.mz=mz
    return

  def get(self,n):
    k= n % len(self.mz)
    n0 = n / len(self.mz)
    return self.g( self.mz[k] , n0 )

  def f(self,_x):
    x=list(_x)
    m=[]
    while True:
      if all(i==0 for i in x):break
      m.append(x[0])
      x=[ x[i+1]-x[i] for i in xrange(len(x)-1) ]
    return m

  def g(self, m, n):
    s=1
    ret=0
    for i,v in enumerate(m):
      ret += v*s
      s = s*(n-i)/(i+1)
    return ret

def f(x,n):
  est=Est()
  est.build(x)
  return est.get(n)

def test():
  z=[(i**3,i+1,i-1)[i%3] for i in xrange(20)]
  print f(z, 300) # 27000000 = 30**3
  print f(z,301) # 301+1
  print f(z,302) # 302-1

test()