Lagrange Polynomial

import numpy as np
import matplotlib.pyplot as plt


class LagrangePolynomial:

    def __init__(self, x, y):
        self.k = len(x)
        self.x = np.array(x)
        self.y = np.array(y)

    def basis(self, x_, j):
        b = [(x_ - self.x[m]) / (self.x[j] - self.x[m]) for m in range(self.k) if m != j]
        return np.prod(b, axis=0) * self.y[j]

    def interpolate(self, x_):
        b = [self.basis(x_, j) for j in range(self.k)]
        return np.sum(b, axis=0)


def main():
    x = [-9, -4, -1, 7]
    y = [5, 2, -2, 9]

    plt.scatter(x, y, c='k')

    lp = LagrangePolynomial(x, y)

    x_ = np.arange(-100, 100) / 10

    plt.plot(x_, lp.basis(x_, 0))
    plt.plot(x_, lp.basis(x_, 1))
    plt.plot(x_, lp.basis(x_, 2))
    plt.plot(x_, lp.basis(x_, 3))
    plt.plot(x_, lp.interpolate(x_), linestyle=':')

    plt.show()


if __name__ == '__main__':
    main()