Kitamasa's Algorithm
(math/kitamasa.cpp)

• View this file on GitHub
• Last update: 2022-03-31 14:37:51+09:00

Description

Kitamasa 法は，$d$ 階線形漸化式の第 $n$ 項を高速に求めるアルゴリズムである．

Operations

• T kitamasa(vector<T> a, vector<T> c, long long n)
  • 初めの $d$ 項 $a_0, \dots, a_{d-1}$ と漸化式 $a_n = c_0 a_{n-d} + \dots + c_{d-1} a_{n-1}$ によって定まる数列 $(a_n)$ の第 $n$ 項を求める．
  • 時間計算量: $O(d^2 \log n)$

Reference

• きたまさ法メモ

Code

#pragma once
#include <vector>

template <typename T>
T kitamasa(const std::vector<T>& a, const std::vector<T>& c, long long n) {
    const int d = a.size();
    if (n < d) {
        return a[n];
    }

    auto dfs = [&](const auto& dfs, long long n) -> std::vector<T> {
        if (n == d) {
            return c;
        }
        std::vector<T> res(d);
        if (n & 1 || n < 2 * d) {
            auto x = dfs(dfs, n - 1);
            for (int j = 0; j < d; ++j) {
                res[j] = c[j] * x[d - 1];
                if (j > 0) {
                    res[j] += x[j - 1];
                }
            }
        } else {
            std::vector x(d, std::vector<T>(d));
            x[0] = dfs(dfs, n / 2);
            for (int i = 1; i < d; ++i) {
                for (int j = 0; j < d; ++j) {
                    x[i][j] = c[j] * x[i - 1][d - 1];
                    if (j > 0) {
                        x[i][j] += x[i - 1][j - 1];
                    }
                }
            }
            for (int i = 0; i < d; ++i) {
                for (int j = 0; j < d; ++j) {
                    res[j] += x[0][i] * x[i][j];
                }
            }
        }
        return res;
    };

    auto x = dfs(dfs, n);
    T ans = 0;
    for (int i = 0; i < d; ++i) {
        ans += x[i] * a[i];
    }
    return ans;
}

#line 2 "math/kitamasa.cpp"
#include <vector>

template <typename T>
T kitamasa(const std::vector<T>& a, const std::vector<T>& c, long long n) {
    const int d = a.size();
    if (n < d) {
        return a[n];
    }

    auto dfs = [&](const auto& dfs, long long n) -> std::vector<T> {
        if (n == d) {
            return c;
        }
        std::vector<T> res(d);
        if (n & 1 || n < 2 * d) {
            auto x = dfs(dfs, n - 1);
            for (int j = 0; j < d; ++j) {
                res[j] = c[j] * x[d - 1];
                if (j > 0) {
                    res[j] += x[j - 1];
                }
            }
        } else {
            std::vector x(d, std::vector<T>(d));
            x[0] = dfs(dfs, n / 2);
            for (int i = 1; i < d; ++i) {
                for (int j = 0; j < d; ++j) {
                    x[i][j] = c[j] * x[i - 1][d - 1];
                    if (j > 0) {
                        x[i][j] += x[i - 1][j - 1];
                    }
                }
            }
            for (int i = 0; i < d; ++i) {
                for (int j = 0; j < d; ++j) {
                    res[j] += x[0][i] * x[i][j];
                }
            }
        }
        return res;
    };

    auto x = dfs(dfs, n);
    T ans = 0;
    for (int i = 0; i < d; ++i) {
        ans += x[i] * a[i];
    }
    return ans;
}

Back to top page