code/math/linear.cc

View this file on GitHub
Last update: 2021-01-23 14:30:21+01:00

Depends on

code/template.cc

Code

#include "../template.hh"
template<typename T, int R, int C>
struct matrix : public array<array<T, C>, R> {
  matrix() : array<array<T, C>, R>{} {}
  matrix(const T& val) : array<array<T, C>, R>{} {
    F0R (i, min(R, C)) (*this)[i][i] = val;
  }
  matrix(initializer_list<T> l) : array<array<T, C>, R>{} {
    int i = 0; for (T e : l) (*this)[i++][0] = e;
  }
  matrix(initializer_list<initializer_list<T>> l) : array<array<T, C>, R>{} {
    int i = 0; for (auto& lp : l) copy(ALL(lp), (*this)[i++].begin());
  }
  T& operator()(size_t i) { return (*this)[i][0]; }
  matrix& operator+=(const matrix& o) {
    F0R (r, R) F0R (c, C) (*this)[r][c] += o[r][c];
    return *this;
  }
  matrix& operator-=(const matrix& o) {
    F0R (r, R) F0R (c, C) (*this)[r][c] -= o[r][c];
    return *this;
  }
  matrix& operator*=(T v) {
    F0R (r, R) F0R (c, C) (*this)[r][c] *= v;
    return *this;
  }
  matrix operator-() {
    matrix res{*this};
    F0R (r, R) F0R (c, C) res[r][c] = -res[r][c];
    return res;
  }
  matrix operator+() { return *this; }
  friend matrix operator+(matrix a, const matrix& b) { return a += b; }
  friend matrix operator-(matrix a, const matrix& b) { return a -= b; }
  friend matrix operator*(matrix a, T v) { return a *= v; }
  friend matrix operator*(T v, matrix a) { return a *= v; }
};
template<typename T, int R, int K, int C>
matrix<T, R, C> operator*(const matrix<T, R, K>& a, const matrix<T, K, C>& b) {
  matrix<T, R, C> res{};
  F0R (r, R) F0R (c, C) F0R (k, K) res[r][c] += a[r][k] * b[k][c];
  return res;
}
template<typename T, int N>
using sqmat = matrix<T, N, N>;
template<typename T, int N>
using vec = matrix<T, N, 1>;
template<typename T, int N>
struct linear {
  using M = sqmat<T, N>; M m;
  using V = vec<T, N>; V v;
  linear() : m(1), v() {}
  linear(const M& _m) : m{_m}, v() {}
  linear(const V& _v) : m(1), v{_v} {}
  linear(const M& _m, const V& _v) : m{_m}, v{_v} {}
  V operator()(const V& x) { return calc(x); }
  V calc(const V& x) { return m * x + v; }
  linear combine(const linear& o) { // o(this(x))
    return linear(o.m * m, o.m * v + o.v);
  }
  linear combine(const M& o) {
    return linear(o * m, o * v);
  }
  linear combine(const V& o) {
    return linear(m, v + o);
  }
  linear combine(const M& om, const V& ov) {
    return combine(linear(om, ov));
  }
};
using L = linear<ll, 2>;

#line 1 "code/template.cc"
// this line is here for a reason
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<vi> vvi;
typedef vector<vii> vvii;
#define fi first
#define se second
#define eb emplace_back
#define pb push_back
#define mp make_pair
#define mt make_tuple
#define endl '\n'
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()
#define SZ(x) (int)(x).size()
#define FOR(a, b, c) for (auto a = (b); (a) < (c); ++(a))
#define F0R(a, b) FOR (a, 0, (b))
template <typename T>
bool ckmin(T& a, const T& b) { return a > b ? a = b, true : false; }
template <typename T>
bool ckmax(T& a, const T& b) { return a < b ? a = b, true : false; }
#ifndef DEBUG
#define DEBUG 0
#endif
#define dout if (DEBUG) cerr
#define dvar(...) " [" << #__VA_ARGS__ ": " << (__VA_ARGS__) << "] "
#line 2 "code/math/linear.cc"
template<typename T, int R, int C>
struct matrix : public array<array<T, C>, R> {
  matrix() : array<array<T, C>, R>{} {}
  matrix(const T& val) : array<array<T, C>, R>{} {
    F0R (i, min(R, C)) (*this)[i][i] = val;
  }
  matrix(initializer_list<T> l) : array<array<T, C>, R>{} {
    int i = 0; for (T e : l) (*this)[i++][0] = e;
  }
  matrix(initializer_list<initializer_list<T>> l) : array<array<T, C>, R>{} {
    int i = 0; for (auto& lp : l) copy(ALL(lp), (*this)[i++].begin());
  }
  T& operator()(size_t i) { return (*this)[i][0]; }
  matrix& operator+=(const matrix& o) {
    F0R (r, R) F0R (c, C) (*this)[r][c] += o[r][c];
    return *this;
  }
  matrix& operator-=(const matrix& o) {
    F0R (r, R) F0R (c, C) (*this)[r][c] -= o[r][c];
    return *this;
  }
  matrix& operator*=(T v) {
    F0R (r, R) F0R (c, C) (*this)[r][c] *= v;
    return *this;
  }
  matrix operator-() {
    matrix res{*this};
    F0R (r, R) F0R (c, C) res[r][c] = -res[r][c];
    return res;
  }
  matrix operator+() { return *this; }
  friend matrix operator+(matrix a, const matrix& b) { return a += b; }
  friend matrix operator-(matrix a, const matrix& b) { return a -= b; }
  friend matrix operator*(matrix a, T v) { return a *= v; }
  friend matrix operator*(T v, matrix a) { return a *= v; }
};
template<typename T, int R, int K, int C>
matrix<T, R, C> operator*(const matrix<T, R, K>& a, const matrix<T, K, C>& b) {
  matrix<T, R, C> res{};
  F0R (r, R) F0R (c, C) F0R (k, K) res[r][c] += a[r][k] * b[k][c];
  return res;
}
template<typename T, int N>
using sqmat = matrix<T, N, N>;
template<typename T, int N>
using vec = matrix<T, N, 1>;
template<typename T, int N>
struct linear {
  using M = sqmat<T, N>; M m;
  using V = vec<T, N>; V v;
  linear() : m(1), v() {}
  linear(const M& _m) : m{_m}, v() {}
  linear(const V& _v) : m(1), v{_v} {}
  linear(const M& _m, const V& _v) : m{_m}, v{_v} {}
  V operator()(const V& x) { return calc(x); }
  V calc(const V& x) { return m * x + v; }
  linear combine(const linear& o) { // o(this(x))
    return linear(o.m * m, o.m * v + o.v);
  }
  linear combine(const M& o) {
    return linear(o * m, o * v);
  }
  linear combine(const V& o) {
    return linear(m, v + o);
  }
  linear combine(const M& om, const V& ov) {
    return combine(linear(om, ov));
  }
};
using L = linear<ll, 2>;