Python 从一元多项式中提取系数和次数，并进行简单的运算

import re    # 正则
from collections import defaultdict    # defaultdict



# 提取一元多项式(type: str)中的次数和系数并转化为字典 -> {次数：系数}
# 例如：'ax^b' -> {b: a}, '-cx^d' -> {d: -c}
# 注意！输入需要是类似这种形式：2x^3 + 1 或者 2x^3+1x^0，不能是 2*x**3 + 1 或者 2*x^3 + 1
def poly_to_dict(polynomial):
    polynomial = polynomial.replace(' ', '')
    dictionary = defaultdict(int)
    positions = [0, *(idx for idx, val in enumerate(polynomial) if val == '+' or val == '-'), len(polynomial)]
    # 一个多项式是由多个单项式组成
    # 而每个单项式都有5个部分：符号 + 系数 + 元 + 脱字符 + 次数
    # 所以匹配规则需要有5个部分，其中“符号，系数，元，次数”有用
    pattern = '([\+-])?(\d+)?([a-z])?\^?(\d+)?'
    for i in range(len(positions) - 1):
        monomial = polynomial[positions[i]: positions[i+1]]
        sign, coeff, var, power = re.match(pattern, monomial).groups()
        # 不要学这种if，if-else写法，我只是为了图方便好看才这样写
        if not sign : sign  = '+'
        if not coeff: coeff = '1' if var else '0'
        if not power: power = '1' if var else '0'
        dictionary[int(power)] += int(sign + coeff)
    return dictionary


# （提取次数和系数后）将其进行简单的运算，比如将两个一元多项式相加
# 其实就是将两个字典中key相同的把value加起来
def add_of_poly(dict1, dict2):
    dict3 = dict1.copy()
    for k2, v2 in dict2.items():
        dict3[k2] += v2
    return {k3: dict3[k3] for k3 in sorted(dict3, reverse = True) if dict3[k3]}


# （提取次数和系数后）将其进行简单的运算，比如将两个一元多项式相乘
# 其实就是将两个字典中key加起来，并把value相乘
def mult_of_poly(dict1, dict2):
    dict3 = defaultdict(int)
    for k1, v1 in dict1.items():
        for k2, v2 in dict2.items():
            dict3[k1 + k2] += v1 * v2
    return {k3: dict3[k3] for k3 in sorted(dict3, reverse = True) if dict3[k3]}


# 将字典形式的多项式按照次数从大到小的顺序打印出来
# 比如：-2x^3 + 5 或者 3x^6 - 2x + 4 这种形式
def print_poly(dictionary):
    output = ''
    for power, coeff in dictionary.items():
        symbol = '+' if coeff > 0 else '-'
        variable = 'x' if power != 0 else ''
        caret, exponent = ('^', power) if power != 0 and power != 1 else ('', '')
        # 5个组成部分
        output += ' {} {}{}{}{}'.format(symbol, abs(coeff), variable, caret, exponent)
    # 将最终输出再次修正一下
    if output:
        if output[3] == '1': output = output[:3] + output[4:]
        print(output[3:] if output[1] == '+' else output[1] + output[3:])
    else:
        print(0)






# 这个函数用作测试，可以看到只要输入满足+ax^b这种形式，就可以正确地得出结果
# 尽管在数学角度看，输入的多项式中有些是不合理的
def test():
    '''
    >>> a = '     2x^3   + x + 5  -                 1x^0'
    >>> b = '-  6x^2 - x - 2x -x   +   x - 5x^6 + 0 - 0x^3  '
    >>> dict_a = poly_to_dict(a)
    >>> dict_b = poly_to_dict(b)
    >>> print_poly(add_of_poly(dict_a, dict_b))
    -5x^6 + 2x^3 - 6x^2 - 2x + 4


    >>> a = ''
    >>> b = 'x'
    >>> c = '-x'
    >>> dict_a = poly_to_dict(a)
    >>> dict_b = poly_to_dict(b)
    >>> dict_c = poly_to_dict(c)

    >>> print_poly(add_of_poly(dict_a, dict_b))
    x
    >>> print_poly(add_of_poly(dict_a, dict_c))
    -x
    >>> print_poly(add_of_poly(dict_b, dict_c))
    0

    >>> print_poly(mult_of_poly(dict_a, dict_a))
    0
    >>> print_poly(mult_of_poly(dict_a, dict_b))
    0
    >>> print_poly(mult_of_poly(dict_a, dict_c))
    0
    >>> print_poly(mult_of_poly(dict_b, dict_a))
    0
    >>> print_poly(mult_of_poly(dict_b, dict_b))
    x^2
    >>> print_poly(mult_of_poly(dict_b, dict_c))
    -x^2
    >>> print_poly(mult_of_poly(dict_c, dict_a))
    0
    >>> print_poly(mult_of_poly(dict_c, dict_b))
    -x^2
    >>> print_poly(mult_of_poly(dict_c, dict_c))
    x^2
    '''


# 引入doctest进行测试
if __name__ == '__main__':
    import doctest
    doctest.testmod()