Compute derivative polinomial

polynomial.py

# coding: utf-8

# Copyright © 2017 Antoine Catton
# This work is free. You can redistribute it and/or modify it under the
# terms of the Do What The Fuck You Want To Public License, Version 2,
# as published by Sam Hocevar. See http://www.wtfpl.net/ for more details.

"""This is inspired from some mythical interview question (Google/Amazon/Microsoft, who knows..)

The goal is "compute the derivative of a polynomial".
"""

import collections
import itertools
import operator

T = collections.namedtuple('Term', 'factor power')

def parse(s):
    def parse_term(t):
        factor_string, sep, power_string = t.partition('x^')
        factor = int(factor_string.strip() or '1')
        power = int(power_string.strip()) if sep else 0
        return T(factor, power)

    return [parse_term(t.strip()) for t in s.split('+')]


assert parse('x^3') == [T(factor=1, power=3)]
assert parse('3x^2') == [T(factor=3, power=2)]
assert parse('4x^2 + 5x^2 + 8x^17') == [T(4, 2), T(5, 2), T(8, 17)]
assert parse('1 + 2x^1') == [T(factor=1, power=0), T(factor=2, power=1)]


def normalize(polynomial):
    power = operator.attrgetter('power')
    ret = (
        T(factor=sum(t.factor for t in terms), power=p)
        for p, terms in itertools.groupby(sorted(polynomial, key=power), key=power)
    )
    return [t for t in ret if t.factor != 0]

assert normalize(parse('2x^3 + 4x^5 + 6x^3')) == [T(8, 3), T(4, 5)]
assert normalize(parse('2x^3 + -2x^3')) == []


def derivative(polynomial):
    assert not any(t.power < 0 for t in polynomial), "This algorithm doesn't support negative powers"
    return [T(factor=t.factor * t.power, power=t.power - 1) for t in polynomial
            if t.power != 0]


assert derivative(parse('x^2')) == parse('2x^1')
assert derivative(parse('1 + 2x^1 + 3x^2')) == parse('2 + 6x^1')