fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. typedef long long ll;
  5. const ll MOD = 1'000'000'007;   // can be changed per problem
  6.  
  7. // ===============================
  8. // 1) Matrix definition and basic operations
  9. // ===============================
  10.  
  11. /**
  12.  * Struct: Matrix
  13.  * Purpose: Represents a matrix of size n x m with long long entries.
  14.  * Usage: Matrix M(n, m); or Matrix M(n, m, data); where data is a 2D vector.
  15.  * Notes: All arithmetic is done modulo MOD.
  16.  *        - The multiplication operator (*) performs matrix multiplication.
  17.  *        - The operator* with vector multiplies matrix by a column vector.
  18.  *        - identity(size) creates an identity matrix.
  19.  * Constraints: n, m >= 0. For multiplication, dimensions must agree.
  20.  * Time Complexity:
  21.  *   - Constructor: O(n*m)
  22.  *   - identity: O(size^2)
  23.  *   - Multiplication (Matrix*Matrix): O(n * m * other.m) with skip-if-zero optimisation.
  24.  *   - Multiplication (Matrix*Vector): O(n * m)
  25.  */
  26. struct Matrix {
  27.     int n, m;                 // dimensions (n rows, m columns)
  28.     vector<vector<ll>> a;     // data
  29.  
  30.     Matrix(int n_ = 0, int m_ = 0) : n(n_), m(m_) {
  31.         a.assign(n, vector<ll>(m, 0));
  32.     }
  33.  
  34.     Matrix(int n_, int m_, const vector<vector<ll>>& data) : n(n_), m(m_), a(data) {}
  35.  
  36.     // Creates an identity matrix (square)
  37.     static Matrix identity(int size) {
  38.         Matrix I(size, size);
  39.         for (int i = 0; i < size; ++i) I.a[i][i] = 1;
  40.         return I;
  41.     }
  42.  
  43.     // Matrix multiplication (with modulo)
  44.     Matrix operator*(const Matrix& other) const {
  45.         if (m != other.n) throw invalid_argument("Incompatible dimensions for multiplication");
  46.         Matrix res(n, other.m);
  47.         for (int i = 0; i < n; ++i) {
  48.             for (int k = 0; k < m; ++k) {
  49.                 if (a[i][k] == 0) continue;      // small optimisation
  50.                 for (int j = 0; j < other.m; ++j) {
  51.                     res.a[i][j] = (res.a[i][j] + a[i][k] * other.a[k][j]) % MOD;
  52.                 }
  53.             }
  54.         }
  55.         return res;
  56.     }
  57.  
  58.     // Multiply matrix by a column vector
  59.     vector<ll> operator*(const vector<ll>& vec) const {
  60.         if (m != (int)vec.size()) throw invalid_argument("Vector size mismatch");
  61.         vector<ll> res(n, 0);
  62.         for (int i = 0; i < n; ++i) {
  63.             for (int j = 0; j < m; ++j) {
  64.                 res[i] = (res[i] + a[i][j] * vec[j]) % MOD;
  65.             }
  66.         }
  67.         return res;
  68.     }
  69.  
  70.     // Print the matrix (for testing)
  71.     void print() const {
  72.         for (int i = 0; i < n; ++i) {
  73.             for (int j = 0; j < m; ++j) cout << a[i][j] << ' ';
  74.             cout << '\n';
  75.         }
  76.     }
  77. };
  78.  
  79. // ===============================
  80. // 2) Fast exponentiation of matrices
  81. // ===============================
  82.  
  83. /**
  84.  * Function: matPow
  85.  * Purpose: Raise a square matrix to a non‑negative integer exponent (binary exponentiation).
  86.  * Usage: Matrix result = matPow(base, exponent);
  87.  * Time Complexity: O(log(exponent) * n^3) where n is the matrix dimension.
  88.  * Notes: The matrix must be square.
  89.  * Constraints: exponent >= 0.
  90.  */
  91. Matrix matPow(Matrix base, ll exponent) {
  92.     if (base.n != base.m) throw invalid_argument("Matrix must be square");
  93.     Matrix result = Matrix::identity(base.n);
  94.     while (exponent > 0) {
  95.         if (exponent & 1) result = result * base;
  96.         base = base * base;
  97.         exponent >>= 1;
  98.     }
  99.     return result;
  100. }
  101.  
  102. // ===============================
  103. // 3) Direct applications: linear recurrences
  104. // ===============================
  105.  
  106. /**
  107.  * Function: fib
  108.  * Purpose: Compute the n-th Fibonacci number (F_0 = 0, F_1 = 1).
  109.  * Usage: ll ans = fib(n);
  110.  * Time Complexity: O(log n) because matrix is 2x2.
  111.  * Notes: Result is modulo MOD. Uses a transition matrix.
  112.  * Constraints: n >= 0.
  113.  */
  114. ll fib(ll n) {
  115.     if (n == 0) return 0;
  116.     if (n == 1) return 1;
  117.     Matrix T(2, 2);
  118.     T.a = {{1, 1}, {1, 0}};   // transition matrix
  119.     Matrix Tn = matPow(T, n - 1);
  120.     vector<ll> initial = {1, 0};   // [F_1, F_0]^T
  121.     vector<ll> result = Tn * initial;
  122.     return result[0];
  123. }
  124.  
  125. /**
  126.  * Function: tribonacci
  127.  * Purpose: Compute the n-th Tribonacci number (T_0=0, T_1=0, T_2=1).
  128.  * Usage: ll ans = tribonacci(n);
  129.  * Time Complexity: O(log n) with 3x3 matrix.
  130.  * Notes: Modulo MOD.
  131.  * Constraints: n >= 0.
  132.  */
  133. ll tribonacci(ll n) {
  134.     if (n == 0 || n == 1) return 0;
  135.     if (n == 2) return 1;
  136.     Matrix T(3, 3);
  137.     T.a = {{1, 1, 1}, {1, 0, 0}, {0, 1, 0}};
  138.     Matrix Tn = matPow(T, n - 2);
  139.     vector<ll> initial = {1, 0, 0};   // [T_2, T_1, T_0]^T
  140.     vector<ll> res = Tn * initial;
  141.     return res[0];
  142. }
  143.  
  144. /**
  145.  * Function: linearRecurrence
  146.  * Purpose: Compute the n-th term of a general linear recurrence of order k:
  147.  *          f(n) = c[0]*f(n-1) + c[1]*f(n-2) + ... + c[k-1]*f(n-k)
  148.  * Usage: ll ans = linearRecurrence(initial, coeff, n);
  149.  *        - initial: vector of size k with f(0), f(1), ..., f(k-1)
  150.  *        - coeff:   vector of size k with c[0], c[1], ..., c[k-1]
  151.  *        - n:       the index to compute (n >= 0)
  152.  * Time Complexity: O(k^3 log n) because matrix is k x k.
  153.  * Notes: Uses the companion matrix. All computations modulo MOD.
  154.  * Constraints: k >= 1, n >= 0.
  155.  */
  156. ll linearRecurrence(const vector<ll>& initial, const vector<ll>& coeff, ll n) {
  157.     int k = initial.size();
  158.     if (n < k) return initial[n] % MOD;
  159.  
  160.     // Build companion matrix
  161.     Matrix T(k, k);
  162.     for (int j = 0; j < k; ++j) T.a[0][j] = coeff[j] % MOD;
  163.     for (int i = 1; i < k; ++i) {
  164.         T.a[i][i-1] = 1;
  165.     }
  166.  
  167.     Matrix Tn = matPow(T, n - k + 1);
  168.     vector<ll> initVec(k);
  169.     for (int i = 0; i < k; ++i) initVec[i] = initial[i] % MOD;
  170.  
  171.     vector<ll> res = Tn * initVec;
  172.     return res[0];
  173. }
  174.  
  175. // ===============================
  176. // 4) Counting walks in a graph
  177. // ===============================
  178.  
  179. /**
  180.  * Function: countWalks
  181.  * Purpose: Count the number of walks of length k from node s to node t in an unweighted directed graph.
  182.  * Usage: ll ans = countWalks(adj, s, t, k);
  183.  *        - adj: adjacency matrix (n x n) where adj.a[i][j] = 1 if edge i->j exists.
  184.  *        - s,t: 0‑based node indices.
  185.  *        - k:    length of the walk (number of edges).
  186.  * Time Complexity: O(log k * n^3) because we raise the adjacency matrix to power k.
  187.  * Notes: Works for both directed and undirected (if symmetric). Result modulo MOD.
  188.  * Constraints: adj must be square; k >= 0.
  189.  */
  190. ll countWalks(const Matrix& adj, int s, int t, ll k) {
  191.     if (adj.n != adj.m) throw invalid_argument("Adjacency matrix must be square");
  192.     Matrix Ak = matPow(adj, k);
  193.     return Ak.a[s][t] % MOD;
  194. }
  195.  
  196. // ===============================
  197. // 5) Sum of first terms of a linear recurrence
  198. // ===============================
  199.  
  200. /**
  201.  * Function: sumFirstNFib
  202.  * Purpose: Compute S(n) = F_0 + F_1 + ... + F_n (sum of first n+1 Fibonacci numbers).
  203.  * Usage: ll ans = sumFirstNFib(n);
  204.  * Time Complexity: O(log n) with a 3x3 augmented matrix.
  205.  * Notes: n >= 0. Result modulo MOD.
  206.  *        The augmented matrix includes the cumulative sum as a state.
  207.  */
  208. ll sumFirstNFib(ll n) {
  209.     if (n == 0) return 0;
  210.     Matrix T(3, 3);
  211.     T.a = {{1, 1, 0}, {1, 0, 0}, {1, 1, 1}};
  212.     Matrix Tn = matPow(T, n);
  213.     vector<ll> initial = {1, 0, 0};  // F_1=1, F_0=0, S_0=0
  214.     vector<ll> res = Tn * initial;
  215.     return res[2];  // S_n
  216. }
  217.  
  218. // The same idea can be generalised to any linear recurrence by adding a row for the sum.
  219.  
  220. // ===============================
  221. // 6) Advanced improvements and tricks
  222. // ===============================
  223.  
  224. /**
  225.  * Function: vecPow
  226.  * Purpose: Compute (base^exp) * vec efficiently without multiplying matrices together at each step.
  227.  * Usage: vector<ll> result = vecPow(base, vec, exp);
  228.  *        - base: square matrix
  229.  *        - vec:  initial column vector
  230.  *        - exp:  exponent (non‑negative)
  231.  * Time Complexity: O(log exp * n^2) because we multiply matrix by vector (O(n^2)) instead of matrix*matrix (O(n^3)).
  232.  * Notes: This is faster when only the final vector is needed and the vector size is small.
  233.  *        The matrix must be square.
  234.  * Constraints: exp >= 0, vec size = base.n.
  235.  */
  236. vector<ll> vecPow(Matrix base, vector<ll> vec, ll exp) {
  237.     while (exp > 0) {
  238.         if (exp & 1) vec = base * vec;
  239.         base = base * base;
  240.         exp >>= 1;
  241.     }
  242.     return vec;
  243. }
  244.  
  245. // Other tricks (sparse matrices, Kitamasa, Berlekamp‑Massey) are included below.
  246.  
  247. /**
  248.  * Function: combine
  249.  * Purpose: Helper for Kitamasa method: multiply two polynomials modulo the characteristic polynomial.
  250.  * Usage: vector<ll> res = combine(a, b, coeff);
  251.  *        - a, b: polynomials as vectors of coefficients (length k)
  252.  *        - coeff: recurrence coefficients (c[0..k-1]) such that
  253.  *                 x^k = coeff[0]*x^(k-1) + ... + coeff[k-1]
  254.  * Time Complexity: O(k^2)
  255.  * Notes: Internal function used by kitamasa.
  256.  */
  257. vector<ll> combine(const vector<ll>& a, const vector<ll>& b, const vector<ll>& coeff) {
  258.     int k = coeff.size();
  259.     vector<ll> res(2 * k, 0);
  260.     for (int i = 0; i < k; i++)
  261.         for (int j = 0; j < k; j++)
  262.             res[i + j] = (res[i + j] + a[i] * b[j]) % MOD;
  263.  
  264.     for (int i = 2*k - 2; i >= k; i--) {
  265.         for (int j = 1; j <= k; j++)
  266.             res[i - j] = (res[i - j] + res[i] * coeff[j-1]) % MOD;
  267.     }
  268.     res.resize(k);
  269.     return res;
  270. }
  271.  
  272. /**
  273.  * Function: kitamasa
  274.  * Purpose: Compute the n-th term of a linear recurrence using Kitamasa's algorithm (O(k^2 log n)).
  275.  * Usage: ll ans = kitamasa(n, init, coeff);
  276.  *        - n: index to compute (n >= 0)
  277.  *        - init: initial terms f(0)..f(k-1)
  278.  *        - coeff: recurrence coefficients (same order as in linearRecurrence)
  279.  * Time Complexity: O(k^2 log n)
  280.  * Notes: Faster than matrix exponentiation when k is large (e.g., k up to a few thousand).
  281.  *        This method avoids matrix multiplication and works with polynomial exponents.
  282.  * Constraints: k >= 1, n >= 0.
  283.  */
  284. ll kitamasa(ll n, const vector<ll>& init, const vector<ll>& coeff) {
  285.     int k = coeff.size();
  286.     if (n < k) return init[n] % MOD;
  287.  
  288.     vector<ll> pol(k, 0), e(k, 0);
  289.     pol[0] = 1;       // represents x^0
  290.     // initialise e to represent x^1
  291.     if (k == 1) {
  292.         // For order 1, x ≡ coeff[0] (mod x - coeff[0])
  293.         e[0] = coeff[0] % MOD;
  294.     } else {
  295.         e[1] = 1;      // x^1
  296.     }
  297.  
  298.     while (n) {
  299.         if (n & 1) pol = combine(pol, e, coeff);
  300.         e = combine(e, e, coeff);
  301.         n >>= 1;
  302.     }
  303.  
  304.     ll ans = 0;
  305.     for (int i = 0; i < k; i++)
  306.         ans = (ans + pol[i] * init[i]) % MOD;
  307.     return ans;
  308. }
  309.  
  310. // ===============================
  311. // 7) Sparse Matrix Exponentiation using Berlekamp‑Massey + Kitamasa
  312. // ===============================
  313.  
  314. /**
  315.  * Function: modPow
  316.  * Purpose: Fast modular exponentiation (a^e % MOD).
  317.  * Usage: ll result = modPow(a, e);
  318.  * Time Complexity: O(log e)
  319.  */
  320. ll modPow(ll a, ll e) {
  321.     ll res = 1;
  322.     a %= MOD;
  323.     while (e > 0) {
  324.         if (e & 1) res = (res * a) % MOD;
  325.         a = (a * a) % MOD;
  326.         e >>= 1;
  327.     }
  328.     return res;
  329. }
  330.  
  331. /**
  332.  * Function: modInv
  333.  * Purpose: Modular inverse of a modulo MOD (MOD must be prime).
  334.  * Usage: ll inv = modInv(a);
  335.  * Time Complexity: O(log MOD)
  336.  */
  337. ll modInv(ll a) {
  338.     return modPow(a, MOD - 2);
  339. }
  340.  
  341. /**
  342.  * Function: dot
  343.  * Purpose: Dot product of two vectors modulo MOD.
  344.  * Usage: ll val = dot(a, b);
  345.  * Time Complexity: O(n)
  346.  */
  347. ll dot(const vector<ll>& a, const vector<ll>& b) {
  348.     ll res = 0;
  349.     for (size_t i = 0; i < a.size(); i++)
  350.         res = (res + a[i] * b[i]) % MOD;
  351.     return res;
  352. }
  353.  
  354. /**
  355.  * Function: berlekamp_massey
  356.  * Purpose: Given the first terms of a linearly recurrent sequence, find the minimal recurrence coefficients.
  357.  * Usage: vector<ll> coeff = berlekamp_massey(s);
  358.  *        - s: vector of initial sequence values (length at least 2 * expected order)
  359.  * Returns: coefficients [c0, c1, ..., cL] such that s[n] = sum_{j=1..L} coeff[j] * s[n-j]
  360.  * Time Complexity: O(L^2) where L is the order found.
  361.  * Notes: This is the Berlekamp‑Massey algorithm. It works modulo MOD.
  362.  *        The returned vector has length L+1 (with constant term = 1, reversed internally).
  363.  * Constraints: MOD must be prime for modular inverse; s must be long enough.
  364.  */
  365. vector<ll> berlekamp_massey(const vector<ll>& s) {
  366.     vector<ll> C(1, 1), B(1, 1);
  367.     ll b = 1; int L = 0, m = 1;
  368.     for (int n = 0; n < (int)s.size(); n++) {
  369.         ll d = s[n];
  370.         for (int i = 1; i <= L; i++)
  371.             d = (d + C[i] * s[n - i]) % MOD;
  372.         if (d == 0) { m++; continue; }
  373.         vector<ll> T = C;
  374.         ll coef = d * modInv(b) % MOD;   // modInv is now defined
  375.         if (C.size() < B.size() + m) C.resize(B.size() + m, 0);
  376.         for (int j = 0; j < (int)B.size(); j++)
  377.             C[j + m] = (C[j + m] - coef * B[j]) % MOD;
  378.         if (2 * L <= n) {
  379.             L = n + 1 - L;
  380.             B = T;
  381.             b = d;
  382.             m = 1;
  383.         } else m++;
  384.     }
  385.     C.resize(L + 1);
  386.     reverse(C.begin(), C.end()); // now coefficients for recurrence
  387.     return C;
  388. }
  389.  
  390. /**
  391.  * Function: sparseMatPow
  392.  * Purpose: Compute a^T * M^k * b for a sparse matrix M, using Berlekamp‑Massey + Kitamasa.
  393.  * Usage: ll result = sparseMatPow(M, a, b, k);
  394.  *        - M: square sparse matrix (n x n)
  395.  *        - a, b: vectors of size n (column vectors)
  396.  *        - k: exponent (non‑negative)
  397.  * Time Complexity: O(n * nnz + L^2 log k) where nnz is number of non‑zero entries in M,
  398.  *                  and L is the order of the recurrence (≤ n).
  399.  * Notes: This is efficient for large n but sparse M. It first generates the sequence
  400.  *        s[i] = a^T * M^i * b for i=0..2n, finds the linear recurrence via Berlekamp‑Massey,
  401.  *        then computes s[k] using Kitamasa.
  402.  *        NOTE: The current implementation multiplies M by a vector using dense O(n^2)
  403.  *        because Matrix uses dense storage. For true sparse efficiency, you should replace
  404.  *        `cur = M * cur` with a sparse multiplication routine.
  405.  * Constraints: M must be square; n >= 1; k >= 0.
  406.  */
  407. ll sparseMatPow(const Matrix& M, const vector<ll>& a, const vector<ll>& b, ll k) {
  408.     int n = M.n;
  409.     vector<ll> s(2 * n + 1);
  410.     vector<ll> cur = b;
  411.     for (int i = 0; i <= 2*n; i++) {
  412.         s[i] = dot(a, cur);   // dot is now defined
  413.         cur = M * cur;        // O(n^2) – replace with sparse version if needed
  414.     }
  415.     vector<ll> coeff = berlekamp_massey(s); // recurrence coefficients
  416.     // coeff: c0, c1, ..., cL such that s[i] = sum_{j=1..L} coeff[j] * s[i-j]
  417.     return kitamasa(k, s, coeff);
  418. }
  419.  
  420. // ===============================
  421. // 8) Precomputation of matrix powers for multiple queries
  422. // ===============================
  423.  
  424. /**
  425.  * Function: precomputePowers
  426.  * Purpose: Precompute powers of a matrix up to a given maximum exponent (for fast queries).
  427.  * Usage: vector<Matrix> powers = precomputePowers(base, maxExp);
  428.  *        - base: square matrix
  429.  *        - maxExp: maximum exponent to support (exclusive, i.e., 2^(maxExp-1) <= maxExp)
  430.  * Time Complexity: O(log(maxExp) * n^3)
  431.  * Notes: The vector powers[i] = base^(2^i). Used with applyPower.
  432.  */
  433. vector<Matrix> precomputePowers(Matrix base, ll maxExp) {
  434.     vector<Matrix> powers;
  435.     powers.push_back(base);
  436.     for (int i = 1; (1LL << i) <= maxExp; i++)
  437.         powers.push_back(powers.back() * powers.back());
  438.     return powers;
  439. }
  440.  
  441. /**
  442.  * Function: applyPower
  443.  * Purpose: Apply a precomputed power to an initial vector.
  444.  * Usage: vector<ll> result = applyPower(powers, exp, init);
  445.  *        - powers: vector from precomputePowers
  446.  *        - exp:    exponent to apply (non‑negative)
  447.  *        - init:   initial column vector
  448.  * Time Complexity: O(log exp * n^2) because we multiply matrix*vector.
  449.  * Notes: This is faster for many queries with different exponents after one precomputation.
  450.  * Constraints: exp >= 0; exp must fit in the precomputed range.
  451.  */
  452. vector<ll> applyPower(const vector<Matrix>& powers, ll exp, const vector<ll>& init) {
  453.     vector<ll> cur = init;
  454.     int bit = 0;
  455.     while (exp > 0) {
  456.         if (exp & 1) cur = powers[bit] * cur;
  457.         exp >>= 1;
  458.         bit++;
  459.     }
  460.     return cur;
  461. }
  462.  
  463. // ===============================
  464. // 9) Test functions (simple examples)
  465. // ===============================
  466.  
  467. void testFibonacci() {
  468.     cout << "Testing Fibonacci:\n";
  469.     for (int i = 0; i <= 10; i++)
  470.         cout << "F(" << i << ") = " << fib(i) << "\n";
  471.     cout << "\n";
  472. }
  473.  
  474. void testLinearRecurrence() {
  475.     cout << "Testing linear recurrence (Fibonacci):\n";
  476.     vector<ll> init = {0, 1};      // F0=0, F1=1
  477.     vector<ll> coeff = {1, 1};     // F(n) = 1*F(n-1) + 1*F(n-2)
  478.     for (int n = 0; n <= 10; n++)
  479.         cout << "F(" << n << ") = " << linearRecurrence(init, coeff, n) << "\n";
  480.     cout << "\n";
  481. }
  482.  
  483. // ===============================
  484. // 10) Example usage
  485. // ===============================
  486.  
  487. int main() {
  488.     ios::sync_with_stdio(false);
  489.     cin.tie(nullptr);
  490.  
  491.     // Quick tests
  492.     testFibonacci();
  493.     testLinearRecurrence();
  494.  
  495.     // Example: counting walks in a simple graph
  496.     // Graph with edges: 0->1, 1->0, 1->2
  497.     Matrix adj(3, 3);
  498.     adj.a = {{0, 1, 0}, {1, 0, 1}, {0, 1, 0}};
  499.     ll paths = countWalks(adj, 0, 2, 3);  // length 3
  500.     cout << "Walks from 0 to 2 of length 3: " << paths << '\n';
  501.  
  502.     // Example: using vecPow for faster Fibonacci
  503.     Matrix T(2, 2);
  504.     T.a = {{1, 1}, {1, 0}};
  505.     vector<ll> init = {1, 0};  // F_1, F_0
  506.     vector<ll> res = vecPow(T, init, 9);   // computes F_10
  507.     cout << "F_10 via vecPow = " << res[0] << '\n';
  508.  
  509.     return 0;
  510. }
Success #stdin #stdout 0s 5320KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Testing Fibonacci:
F(0) = 0
F(1) = 1
F(2) = 1
F(3) = 2
F(4) = 3
F(5) = 5
F(6) = 8
F(7) = 13
F(8) = 21
F(9) = 34
F(10) = 55

Testing linear recurrence (Fibonacci):
F(0) = 0
F(1) = 1
F(2) = 1
F(3) = 1
F(4) = 2
F(5) = 3
F(6) = 5
F(7) = 8
F(8) = 13
F(9) = 21
F(10) = 34

Walks from 0 to 2 of length 3: 0
F_10 via vecPow = 55
https://ideone.com/hn4RNt
language:
C++14 (clang 8.0)
created:
1 day ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!