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

This commit is contained in:
Raymond Hettinger 2022-08-21 09:41:29 -05:00 committed by GitHub
parent eea8f42ded
commit 0e28a3a50f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 19 additions and 0 deletions

View File

@ -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']