2009-12-06

Polynomials in Python with operator overloading

This blog post is a fun example for implementing univariate polynomials with operator overloading in Python. It is almost a domain-specific language (DSL). Here is the code:

#! /usr/bin/python2.4

class Polynomial(object):

  __attrs__ = ['co']

  def __init__(self, co=None):
    if co:
      co = list(co)
      while len(co) > 1 and not co[-1]:
        co.pop()
      self.co = co
    else:
      self.co = [0]

  def __repr__(self):
    items = []
    for i in xrange(len(self.co)):
      if self.co[i]:
        if i == 0:
          items.append('%s' % self.co[i])
        else:
          if self.co[i] == 1:
            pre = ''
          else:
            pre = '%s * ' % self.co[i]
          if i == 1:
            post = ''
          else:
            post = ' ** %s' % i
          items.append('%sx%s' % (pre, post))
    return ' + '.join(items) or '0'

  __str__ = __repr__

  def __call__(self, x):
    i = len(self.co)
    v = 0
    while i > 0:
      i -= 1
      v = self.co[i] + x * v
    return v

  def __add__(self, p):
    if isinstance(p, Polynomial):
      q = self
      if len(p.co) > len(q.co):
        p, q = q, p
      co = list(q.co)
      for i in xrange(len(p.co)):
        co[i] += p.co[i]
    else:
      co = list(self.co)
      co[0] += p
    return type(self)(co)

  __radd__ = __add__

  def __mul__(self, p):
    if isinstance(p, Polynomial):
      co = []
      for i in xrange(len(self.co) + len(p.co) - 1):
        s = 0
        for j in xrange(max(0, i - len(p.co) + 1), min(i + 1, len(self.co))):
          s += self.co[j] * p.co[i - j]
        co.append(s)
      return type(self)(co)
    else:
      return type(self)([x * p for x in self.co])

  __rmul__ = __mul__

  def __pow__(self, n):
    assert isinstance(n, int)
    assert n >= 0
    if n == 0:
      return type(self)([1])
    co = type(self)(self.co)
    while n > 1:
      co = co * self
      n -= 1
    return co

  #__rpow__ = __pow__

  def __sub__(self, p):
    return self + (-1) * p

  def __rsub__(self, p):
    return (-1) * self + p

x = Polynomial([0, 1])

# --- Test

for p in (x * (x + 1) * (2 * x + 1),
          (x + 1) * (1 + x),
          (x + 2) * (2 + x),
          (x + 3) ** 5,
          0 * x,
          x ** 2 - 2 * x + 1,
          (x + 3) + (3 * x ** 2 + 6) + x ** 4):
  print 'p = ', p
  print 'p(0) = ', p(0)
  print 'p(7) = ', p(7)
  print

Similar code for symbolic calculation with polynomials in Perl.

No comments: