linear_recurrence_relations.cpp

#include <cassert>
#include <iostream>
#include <vector>

template <class T>
using vec = std::vector<T>;

template <class T>
using mat = vec<vec<T>>;

template <class T>
mat<T> operator*(mat<T> const& a, mat<T> const& b) {
  int m = a.size(), n1 = a[0].size(), n2 = b.size(), p = b[0].size();
  assert(n1 == n2);

  mat<T> res(m, vec<T>(p, T(0)));
  for (auto i = 0; i < m; ++i) {
    for (auto j = 0; j < p; ++j) {
      for (auto k = 0; k < n1; ++k) {
        res[i][j] += a[i][k] * b[k][j];
      }
    }
  }

  return res;
}

template <class T>
T binary_exponentiation(T value, uint64_t pow) {
  auto res = value;
  --pow;

  for (; pow > 0; pow >>= 1) {
    if (pow & 1) {
      res = res * value;
    }

    value = value * value;
  }

  return res;
}

auto build_transformation_matrix(vec<int64_t> const& coefficients) {
  auto k = coefficients.size();

  auto result = mat<int64_t>(k, vec<int64_t>(k, 0));
  for (auto i = 0; i < k - 1; ++i) {
    result[i][i + 1] = 1;
  }

  std::copy(coefficients.begin(), coefficients.end(), result[k - 1].begin());

  return result;
}

auto build_values_matrix(vec<int64_t> const& initial_values) {
  auto k = initial_values.size();

  auto result = mat<int64_t>(k, vec<int64_t>(1));
  for (auto i = 0; i < k; ++i) {
    result[i][0] = initial_values[i];
  }

  return result;
}

int64_t solve(uint64_t n,
              vec<int64_t> const& initial_values,
              vec<int64_t> const& coefficients) {
  assert(initial_values.size() == coefficients.size());

  auto k = initial_values.size();
  if (n < k) {
    return initial_values[n];
  }

  auto transformation = build_transformation_matrix(coefficients);
  auto values = build_values_matrix(initial_values);

  auto result = binary_exponentiation(transformation, n) * values;

  return result[0][0];
}