Add polynomial_from_roots() to the itertools recipes (#96102)

rhettinger · web-flow · commit 0e28a3a50ff9 · 2022-08-21T09:41:29.000-05:00
diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
@@ -800,6 +800,18 @@ which incur interpreter overhead.
            window.append(x)
            yield sum(map(operator.mul, kernel, window))
 
+   def polynomial_from_roots(roots):
+       """Compute a polynomial's coefficients from its roots.
+
+          (x - 5) (x + 4) (x - 3)  expands to:   x³ -4x² -17x + 60
+       """
+       # polynomial_from_roots([5, -4, 3]) --> [1, -4, -17, 60]
+       roots = list(map(operator.neg, roots))
+       return [
+           sum(map(math.prod, combinations(roots, k)))
+           for k in range(len(roots) + 1)
+       ]
+
    def flatten(list_of_lists):
        "Flatten one level of nesting"
        return chain.from_iterable(list_of_lists)
@@ -1137,6 +1149,13 @@ which incur interpreter overhead.
     >>> list(convolve(data, [1, -2, 1]))
     [20, 0, -36, 24, -20, 20, -20, -4, 16]
 
+    >>> polynomial_from_roots([5, -4, 3])
+    [1, -4, -17, 60]
+    >>> factored = lambda x: (x - 5) * (x + 4) * (x - 3)
+    >>> expanded = lambda x: x**3 -4*x**2 -17*x + 60
+    >>> all(factored(x) == expanded(x) for x in range(-10, 11))
+    True
+
     >>> list(flatten([('a', 'b'), (), ('c', 'd', 'e'), ('f',), ('g', 'h', 'i')]))
     ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']
 
