Codewars: Binomial Expansion

一道3kyu的题。题目如下：

The purpose of this kata is to write a program that can do some algebra.

Write a function expand that takes in an expresion with a single, one character variable, and expands it. The expresion is in the form (ax+b)^n where a and b are integers which may be positive or negative, x is any one character long variable, and n is a natural number. If a = 1, no coeficient will be placed in front of the variable. If a = -1, a “-” will be placed in front of the variable.

The expanded form should be returned as a string in the form ax^b+cx^d+ex^f… where a, c, and e are the coefficients of the term, x is the original one character variable that was passed in the original expression and b, d, and f, are the powers that x is being raised to in each term and are in decreasing order.

If the coeficient of a term is zero, the term should not be included. If the coeficient of a term is one, the coeficient should not be included. If the coeficient of a term is -1, only the “-” should be included. If the power of the term is 0, only the coeficient should be included. If the power of the term is 1, the carrot and power should be excluded.

---

用正则替换+Python的语法糖能玩出花来。

这里代码比较平凡，主要是一些细节问题。很冗长，就不需要注释了。

---

# version: 3.6
import re
def combination(k, n):
    if k == 0:
        return 1
    c = 1
    if k + 1 < n - k:
        for i in range(n - k + 1, n + 1):
            c *= i
        for i in range(1, k + 1):
            c /= i
    else:
        for i in range(k + 1, n + 1):
            c *= i
        for i in range(1, n - k + 1):
            c /= i
    return c

def quick_pow(base, power):
    res = 1
    while power:
        if power & 1:
            res *= base
            power -= 1
        else:
            base *= base
            power >>= 1
    return res


def term_product(term_left, term_right, pow_left, pow_right, comb):
    def term_pow(term, power):
        num = ''
        coef = 1
        i = 0
        while i < len(term):
            ch = term[i]
            if ch is None: break
            if not ch.isdigit():
                if ch == '+':
                    coef = 1
                elif ch == '-':
                    coef = -1
                else:
                    break
            else:
                num += ch
            i += 1
        if num:
            num = int(quick_pow(int(num), power))
            if power & 1: num = coef * num
        else:
            num = coef if power&1 else 1
        var = term[i:] if power != 0 else ''
        return num, var

    num_left, var_left = term_pow(term_left, pow_left)
    num_right, var_right = term_pow(term_right, pow_right)

    final_num = comb * num_left * num_right

    final_pow = ('', '')
    final_var = (var_left, var_right)
    if var_left and var_right:
        if var_left == var_right:
            final_var = (var_left, '')
            if pow_left + pow_right > 1:
                final_pow = ('^' + str(pow_left + pow_right), '')
        else:
            final_pow = ('^' + str(pow_left), '^' + str(pow_right))
    elif not var_left and var_right:
        if pow_right > 1:
            final_pow = ('', '^' + str(pow_right))
    elif var_left:
        if pow_left > 1:
            final_pow = ('^' + str(pow_left), '')
    if final_var[0] or final_var[1]:
        if final_num==1 or final_num==0: final_num = ''
        elif final_num==-1: final_num = '-'
    if not final_num or final_num!='-' and final_num > 0: final_num = '+' + str(final_num)
    return '{0}{1}{3}{2}{4}'.format(final_num, *final_var, *final_pow)


reg = re.compile(
    r'\((?P<term_left>-?.+)(?P<term_right>[+-].+)\)\^(?P<power>\w+)')


def expand(expr):
    matched = reg.match(expr)
    if not matched:
        return ''
    term_left = matched.group('term_left')
    term_right = matched.group('term_right')
    power = int(matched.group('power'))
    if power == 0:
        return '1'
    expanded = ''
    for k in range(power, -1, -1):
        _comb = int(combination(k, power))
        expanded += term_product(term_left, term_right, k, power - k, _comb)
    return expanded[1:] if expanded[0] == '+' else expanded


class Test:
    @staticmethod
    def assert_equals(expect, true):
        assert expect == true


Test.assert_equals(expand("(x+1)^0"), "1")
Test.assert_equals(expand("(x+1)^1"), "x+1")
Test.assert_equals(expand("(x+1)^2"), "x^2+2x+1")

Test.assert_equals(expand("(x-1)^0"), "1")
Test.assert_equals(expand("(x-1)^1"), "x-1")
Test.assert_equals(expand("(x-1)^2"), "x^2-2x+1")

Test.assert_equals(expand("(5m+3)^4"), "625m^4+1500m^3+1350m^2+540m+81")
Test.assert_equals(expand("(2x-3)^3"), "8x^3-36x^2+54x-27")
Test.assert_equals(expand("(7x-7)^0"), "1")

Test.assert_equals(expand("(-5m+3)^4"), "625m^4-1500m^3+1350m^2-540m+81")
Test.assert_equals(expand("(-2k-3)^3"), "-8k^3-36k^2-54k-27")
Test.assert_equals(expand("(-7x-7)^0"), "1")