github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/math/polynomial.go (about)

     1  package math
     2  
     3  import (
     4  	"fmt"
     5  	"math/big"
     6  	"strings"
     7  )
     8  
     9  // Polynomial is used for educational purpose only, not optimized for sparse polynomials
    10  type Polynomial struct {
    11  	maxOrder     int
    12  	coefficients []*big.Int
    13  }
    14  
    15  func NewPolynomialWithMaxOrder(maxOrder int, coefficients []*big.Int) *Polynomial {
    16  	return &Polynomial{maxOrder: maxOrder, coefficients: coefficients}
    17  }
    18  
    19  func NewPolynomial(coefficients []*big.Int) *Polynomial {
    20  	return &Polynomial{coefficients: coefficients}
    21  }
    22  
    23  var zero = big.NewInt(0)
    24  var one = big.NewInt(1)
    25  
    26  func (p1 *Polynomial) String() string {
    27  	var parts []string
    28  	for i, c := range p1.coefficients {
    29  		if c.Cmp(zero) != 0 {
    30  			if i == 0 {
    31  				parts = append(parts, c.String())
    32  			} else {
    33  				if c.Cmp(one) == 0 {
    34  					parts = append(parts, fmt.Sprintf("x^%d", i))
    35  				} else {
    36  					parts = append(parts, fmt.Sprintf("%sx^%d", c.String(), i))
    37  				}
    38  			}
    39  		}
    40  	}
    41  
    42  	return strings.Join(parts, " + ")
    43  }
    44  
    45  func (p1 *Polynomial) Mul(p2 *Polynomial) *Polynomial {
    46  	dim := len(p1.coefficients) + len(p2.coefficients) - 1
    47  	if dim <= 0 {
    48  		return &Polynomial{}
    49  	}
    50  
    51  	if p1.maxOrder > 0 && dim > p1.maxOrder+1 {
    52  		dim = p1.maxOrder + 1
    53  	}
    54  	coefficients := make([]*big.Int, dim)
    55  	for i, c1 := range p1.coefficients {
    56  		for j, c2 := range p2.coefficients {
    57  			if p1.maxOrder > 0 && p1.maxOrder < i+j {
    58  				break
    59  			}
    60  			c3 := coefficients[i+j]
    61  			if c3 == nil {
    62  				c3 = big.NewInt(0)
    63  			}
    64  
    65  			coefficients[i+j] = c3.Add(c3, big.NewInt(0).Mul(c1, c2))
    66  		}
    67  	}
    68  
    69  	return &Polynomial{maxOrder: p1.maxOrder, coefficients: coefficients}
    70  }