dictで多変数の多項式の同類項をまとめる|Python

多変数の多項式をプログラムで表現すると係数と変数を分けた配列になると思います。 計算過程で同じ変数を持つ項（同類項）をまとめていく方法を考えました。

多変数の多項式

１変数の多項式の場合はnumpy.polynomialを使うと良いです。

Atcoderのマラソン型コンテスト北大・日立新概念コンピューティングコンテスト2018に参加した中で実装しました。

やりたいこと

配列で表現した多項式の同類項をまとめる。

以下は必要になった経緯です。

例

例えば、多項式を競技プログラミングで扱うとき以下のような問題がでます。

入力は以下の形式で標準入力から与えられる。

$\begin{matrix} N & K \\ d_1 & c_1 & v_{1,1} & v_{1,2} & ... & v_{1,d_1} \\ d_2 & c_2 & v_{2,1} & v_{2,2} & ... & v_{2,d_2} \\ . \\ . \\ . \\ d_K & c_K & v_{K,1} & v_{K,2} & ... & v_{K,d_K} \end{matrix}$

---

• 入力はすべて整数である。
• $N$は変数の数、$K$は項の数を表す。
• 2行目から $K+1$行目までの各行では、それぞれの項の情報を表す整数が与えられる。
• $d_i$は第$i$項の次数を、 $c_i$は第$i$項の係数を、$v_{i,1} ... v_{i,d_i}$ はその項を構成する変数番号を表す。
• $d_i = 0$であるときは定数項を表し、その行は$d_i c_i$のみからなる。

atcoderの入力より引用。

以下のような多項式の場合、

これをプログラムで表現すると係数と変数を分けた配列が扱いやすそうです。

C = [1, 3, 4, 5]
V = [[1, 2, 3],
    [1, 4],
    [2],
    []]

ここで、計算の途中で同類項が出現したとき、項をまとめていく方法を考えます。

C = [3, 2, 4, -4, 5]
V = [[1, 2],
    [3],
    [1, 2],
    [3],
    []]

この時、配列をすべて比較して同類項をまとめると面倒くさそうです。

多変数の多項式をまとめる

dictを使って多項式をまとめました。 keyに変数の文字列を使います。

def similar_term(C, V):
    # 同類項をまとめる
    di = {}
    for c, v in zip(C, V):
        if len(v) is 0:
            # 定数項
            di[' '] = c
            continue
        # 文字列にしてKeyにする
        v = ' '.join(map(str, sorted(v)))
        if di.get(v) is None:
            di[v] = c
        else:
            di[v] += c
    # print(di)
    ## {'1 2': 7, '3': -2, ' ': 5}

    # 配列へ戻す
    new_V = [list(map(int, d.split())) for d in di.keys()]
    new_C = list(di.values())
    return new_C, new_V


C = [3, 2, 4, -4, 5]
V = [[1, 2], [3], [1, 2], [3], []]

newC, newV = similar_term(C, V)
print(newC)
print(newV)
# [7, -2, 5]
# [[1, 2], [3], []]

同類項がまとまった。

C = [3, 2, 4, -4, 5]
V = [[2, 1], [1, 2], [], []]

newC, newV = similar_term(C, V)
print(newC)
print(newV)
# [5, -4]
# [[1, 2, 3], []]

定数項や変数の順序が異なってても対応しています。

数式の表示

ここで載せた数式はnpmのkatexを使って変換しています。

```math
E = mc^2
```

こう書くと下のような数式になる。

表記方法はこちらを参照した。