【Python Sympy】将表达式化为关于x的多项式，求出多项式系数

今天遇到一个问题。老师让用C++实现三次样条插值函数，输出多项式系数，其中有一个公式长这样$$S(x)=\frac{m_i\left(x-x_i\right){\left(x-x_{i+1}\right)}^2}{h_i^2}+\frac{m_{i+1}{\left(x-x_i\right)}^2\left(x-x_{i+1}\right)}{h_i^2}+\frac{y_i{\left(x-x_{i+1}\right)}^2\left(h_i+2x-2x_i\right)}{h_i^3}+\frac{y_{i+1}{\left(x-x_i\right)}^2\left(h_i-2x+2x_{i+1}\right)}{h_i^3}$$对应的$\text{LATEX}$代码为

S(x)=\frac{m_{i} \left(x - x_{i}\right) \left(x - x_{i+1}\right)^{2}}{h_{i}^{2}} + \frac{m_{i+1} \left(x - x_{i}\right)^{2} \left(x - x_{i+1}\right)}{h_{i}^{2}} + \frac{y_{i} \left(x - x_{i+1}\right)^{2} \left(h_{i} + 2 x - 2 x_{i}\right)}{h_{i}^{3}} + \frac{y_{i+1} \left(x - x_{i}\right)^{2} \left(h_{i} - 2 x + 2 x_{i+1}\right)}{h_{i}^{3}}

要求出各项系数，就必须把上面这个贼长的式子化成关于x的多项式。……那得累死。怎么办呢？

Python里有一个模块叫sympy，是做符号计算的。安装方法为：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple sympy

sympy里面有一个函数collect，作用是将关于一个/几个变量的同次项合并在一起。函数定义为：

sympy.simplify.radsimp.collect(expr, syms, func=None, evaluate=None, exact=False, distribute_order_term=True)

其中，expr是要化简的表达式，syms指定化简成关于谁的多项式（可以是一个变量或是变量的列表），evaluate表示是否对结果赋值。若evaluate = True，则直接返回化简后的多项式；若evaluate = False，则返回一个字典，存储每一项的系数。

例如，要把$a{x}^2+3x+c+\frac{1}{b}{x}^2+9+\frac{d}{3}x$化为关于$x$的多项式，可以这么写（下面的结果是在IPython中运行得到的）：

In [4]: from sympy import *

In [5]: a, b, c, d, x = symbols('a, b, c, d, x')

In [6]: collect(a * x ** 2 + 3 * x + c + 1 / b * x ** 2 + 9 + d / 3 * x, x)
Out[6]: c + x**2*(a + 1/b) + x*(d/3 + 3) + 9

结果是$\left(a+\frac{1}{b}\right)x^2+\left(\frac{d}{3}+3\right)x+c+9$。

若设evaluate = False：

In [8]: p = collect(a * x ** 2 + 3 * x + c + 1 / b * x ** 2 + 9 + d / 3 * x, x, evaluate = False)

In [9]: p
Out[9]: {x: d/3 + 3, x**2: a + 1/b, 1: c + 9}

In [10]: p[x ** 2] # x平方项的系数
Out[10]: a + 1/b

In [11]: p[1] # 常数项
Out[11]: c + 9

也可以用它来化简多变量的表达式。

In [12]: y = Symbol('y')

In [13]: collect(x ** 2 + y * x ** 2 + x * y + y + a * y, [x, y])
Out[13]: x**2*(y + 1) + x*y + y*(a + 1)

In [14]: collect(x ** 2 + y * x ** 2 + x * y + y + a * y, [x, y], evaluate = False)
Out[14]: {y: a + 1, x**2: y + 1, x: y}

注意这个算法是贪心的，它同一时刻只处理一个变量。

还有更高级的用法：

In [15]: collect(a * sin(2 * x) + b * sin(2 * x), sin(2 * x))
Out[15]: (a + b)*sin(2*x)

In [16]: collect(a * x * log(x) + b * (x * log(x)) + x, x * log(x))
Out[16]: x*(a + b)*log(x) + x

---

现在解决我们一开始提出的问题。注意，要先调用expand展开括号中的东西才能继续化简。

In [18]: S, h_i, x_i, x_j, y_i, y_j, m_i, m_j = symbols('S, h_i, x_i, x_j, y_i, y_j, m_i, m_j')

In [19]: S = (h_i + 2 * (x - x_i)) * (x - x_j) ** 2 * y_i / h_i ** 3 + \
    ...: (h_i - 2 * (x - x_j)) * (x - x_i) ** 2 * y_j / h_i ** 3 + \
    ...: (x - x_i) * (x - x_j) ** 2 * m_i / h_i ** 2 + \
    ...: (x - x_j) * (x - x_i) ** 2 * m_j / h_i ** 2

In [20]: collect(expand(S), x)
Out[20]: x**3*(m_i/h_i**2 + m_j/h_i**2 + 2*y_i/h_i**3 - 2*y_j/h_i**3) + x**2*(-m_i*x_i/h_i**2 - 2*m_i*x_j/h_i**2 - 2*m_j*x_i/h_i**2 - m_j*x_j/h_i**2 + y_i/h_i**2 + y_j/h_i**2 - 2*x_i*y_i/h_i**3 + 4*x_i*y_j/h_i**3 - 4*x_j*y_i/h_i**3 + 2*x_j*y_j/h_i**3) + x*(2*m_i*x_i*x_j/h_i**2 + m_i*x_j**2/h_i**2 + m_j*x_i**2/h_i**2 + 2*m_j*x_i*x_j/h_i**2 - 2*x_i*y_j/h_i**2 - 2*x_j*y_i/h_i**2 - 2*x_i**2*y_j/h_i**3 + 4*x_i*x_j*y_i/h_i**3 - 4*x_i*x_j*y_j/h_i**3 + 2*x_j**2*y_i/h_i**3) - m_i*x_i*x_j**2/h_i**2 - m_j*x_i**2*x_j/h_i**2 + x_i**2*y_j/h_i**2 + x_j**2*y_i/h_i**2 + 2*x_i**2*x_j*y_j/h_i**3 - 2*x_i*x_j**2*y_i/h_i**3

稍稍化简之后得$$S\left(x\right)=\frac{h_i\left(m_i+m_{i+1}\right)+2y_i-2y_{i+1}}{h_i^3}{x}^3+\frac{h_i\left(-m_i{x}_i-2m_i{x}_{i+1}-2m_{i+1}{x}_i-m_{i+1}{x}_{i+1}+y_i+y_{i+1}\right)-2x_i{y}_i+4x_i{y}_{i+1}-4x_{i+1}{y}_i+2x_{i+1}{y}_{i+1}}{h_i^3}{x}^2+\frac{h_i\left(2m_i{x}_i{x}_{i+1}+m_i{x}_{i+1}^2+m_{i+1}{x}_i^2+2m_{i+1}{x}_i{x}_{i+1}-2x_i{y}_{i+1}-2x_{i+1}{y}_i\right)-2x_i^2{y}_{i+1}+4x_i{x}_{i+1}{y}_i-4x_i{x}_{i+1}{y}_{i+1}+2x_{i+1}^2{y}_i}{h_i^3}x+\frac{h_i\left(-m_i{x}_i{x}_{i+1}^2-m_{i+1}{x}_i^2{x}_{i+1}+x_i^2{y}_{i+1}+x_{i+1}^2{y}_i\right)+2x_i{x}_{i+1}\left(x_i{y}_{i+1}-x_{i+1}{y}_i\right)}{h_i^3}$$

对应的$\text{LATEX}$代码为

S\left(x\right)=\frac{h_i\left(m_i+m_{i+1}\right)+2y_i-2y_{i+1}}{h_i^3}x^3+\frac{h_i\left(-m_ix_i-2m_ix_{i+1}-2m_{i+1}x_i-m_{i+1}x_{i+1}+y_i+y_{i+1}\right)-2x_iy_i+4x_iy_{i+1}-4x_{i+1}y_i+2x_{i+1}y_{i+1}}{h_i^3}x^2+\frac{h_i\left(2m_ix_ix_{i+1}+m_ix_{i+1}^2+m_{i+1}x_i^2+2m_{i+1}x_ix_{i+1}-2x_iy_{i+1}-2x_{i+1}y_i\right)-2x_i^2y_{i+1}+4x_ix_{i+1}y_i-4x_ix_{i+1}y_{i+1}+2x_{i+1}^2y_i}{h_i^3}x+\frac{h_i\left(-m_ix_ix_{i+1}^2-m_{i+1}x_i^2x_{i+1}+x_i^2y_{i+1}+x_{i+1}^2y_i\right)+2x_ix_{i+1}\left(x_iy_{i+1}-x_{i+1}y_i\right)}{h_i^3}

至此，问题圆满解决。

---

参考资料

1. https://docs.sympy.org/latest/modules/simplify/simplify.html
2. https://www.geeksforgeeks.org/python-sympy-collect-method/