import numpy as np


# Tri Diagonal Matrix Algorithm(a.k.a. Thomas algorithm) solver
def tdma_solver(a, b, c, d):
    """
    TDMA solver, a b c d can be NumPy array type or Python list type.
    refer to https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
    and to https://www.cfd-online.com/Wiki/Tridiagonal_matrix_algorithm_-_TDMA_(Thomas_algorithm)
    """
    nf = len(d)  # number of equations
    ac, bc, cc, dc = map(np.array, (a, b, c, d))  # copy arrays
    for it in range(1, nf):
        mc = ac[it - 1] / bc[it - 1]
        bc[it] = bc[it] - mc * cc[it - 1]
        dc[it] = dc[it] - mc * dc[it - 1]

    xc = bc
    xc[-1] = dc[-1] / bc[-1]

    for il in range(nf - 2, -1, -1):
        xc[il] = (dc[il] - cc[il] * xc[il + 1]) / bc[il]

    return xc
