# -*- coding: utf-8 -*-
"""
Created on Mon Mar 19 20:18:35 2018

@author: Arseny Mannerheim
"""

import numpy as np

print ("Метод прогонки\n")

A = np.array([[2.1, 0.5, 0, 0, 0], [0.5, 6.1, 0.5, 0, 0], [0, 0.5, 2.5, 0.5, 0], [0, 0, 0.5, 3.4, 0.5], [0, 0, 0, 0.5, 7.15]],dtype=float)

a = np.array([0.5, 0.5, 0.5, 0.5])
b = np.array([2.1, 6.1, 2.5, 3.4, 7.15])
c = np.array([0, 0.5, 0.5, 0.5])
d = np.array([2.6, 7.0, 4.1, 4.5, 6.9])

def TDMAsolver(a, b, c, d):
    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

print ("A:\n", A)
print ("\nB:\n", d)
print ("\nX:\n", TDMAsolver (a, b, c, d))