An example tridiagonal matrix solving routine

Listed below is an example tridiagonal matrix solving routine which utilizes the Blitz++ library (see Sect. 2.20).

// Tridiagonal.cpp

// Function to invert tridiagonal matrix equation.
// Left, centre, and right diagonal elements of matrix
// stored in arrays a, b, c, respectively.
// Right-hand side stored in array w.
// Solution written to array u.

// Matrix is NxN. Arrays a, b, c, w, u assumed to be of extent N+2,
// with redundant 0 and N+1 elements.

#include <blitz/array.h>

using namespace blitz;

void Tridiagonal (Array<double,1> a, Array<double,1> b, Array<double,1> c, 
                  Array<double,1> w, Array<double,1>& u)
{
  // Find N. Declare local arrays.	
  int N = a.extent(0) - 2;
  Array<double,1> x(N), y(N);

  // Scan up diagonal from i = N to 1
  x(N-1) = - a(N) / b(N);
  y(N-1) = w(N) / b(N);
  for (int i = N-2; i > 0; i--)
    {
      x(i) = - a(i+1) / (b(i+1) + c(i+1) * x(i+1));
      y(i) = (w(i+1) - c(i+1) * y(i+1)) / (b(i+1) + c(i+1) * x(i+1));
    }
  x(0) = 0.;
  y(0) = (w(1) - c(1) * y(1)) / (b(1) + c(1) * x(1));

  // Scan down diagonal from i = 0 to N-1
  u(1) = y(0);
  for (int i = 1; i < N; i++)
    u(i+1) = x(i) * u(i) + y(i);
}