Untitled

 avatar
unknown
python
a year ago
785 B
5
Indexable
import numpy as np

def tridiagonal(A, d):
    # Determines n
    n = len(d)

    # Preallocates all necessary vectors
    a = np.zeros(n-1)
    b = np.zeros(n)
    c = np.zeros(n-1)
    x = np.zeros(n)

    # Extracts first element of "b" from "A"
    b[0] = A[0, 0]
    
    # Forward loop
    for i in range(1, n):
        # Extract relevant elements of "a", "b", and "c" from "A"
        a[i-1] = A[i, i-1]
        b[i] = A[i, i]
        c[i-1] = A[i-1, i]

        # Forward elimination
        w = a[i-1] / b[i-1]
        b[i] = b[i] - w * c[i-1]
        d[i] = d[i] - w * d[i-1]

    # Backward loop (backward substitution)
    x[n-1] = d[n-1] / b[n-1]
    for i in range(n-2, -1, -1):
        x[i] = (d[i] - c[i] * x[i+1]) / b[i]

    return x
Editor is loading...
Leave a Comment