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()