Lagrange.cpp

• View this file on GitHub
• Last update: 2023-03-06 11:24:14-03:00

Depends on

• BitTricks.cpp
• Combinatorics.cpp
• bits/stdc++.h

Required by

• Exponential Sum (polynomial/ExponentialSum.cpp)

Verified with

• tests/yosupo/exp-sum-limit.test.cpp
• tests/yosupo/exp-sum.test.cpp

Code

#ifndef _LIB_LAGRANGE
#define _LIB_LAGRANGE
#include <bits/stdc++.h>
#include "Combinatorics.cpp"

namespace lib {
using namespace std;
namespace linalg {
template <typename Field> struct PrefixLagrange {
  vector<Field> pref, suf;
  PrefixLagrange() {}

  void ensure(int n) {
    int o = pref.size();
    if (n <= o)
      return;
    pref.resize(n), suf.resize(n);
  }

  template <typename T> Field eval(const vector<Field> &v, T x) {
    using C = Combinatorics<Field>;
    assert(!v.empty());
    int d = (int)v.size() - 1;
    if (x <= d)
      return v[x];

    ensure(d + 1);

    Field a = x;
    pref[0] = suf[d] = 1;
    for (T i = 0; i < d; i++)
      pref[i + 1] = pref[i] * a, a -= 1;
    for (T i = d; i; i--)
      suf[i - 1] = suf[i] * a, a += 1;

    Field ans = 0;
    for (int i = 0; i <= d; i++) {
      Field l = pref[i] * suf[i] * C::ifactorial(i) * C::ifactorial(d-i) * v[i];
      if ((d + i) & 1)
        l = -l;
      ans += l;
    }
    return ans;
  }
};

template<typename T, typename U>
T lagrange_iota(const vector<T>& f, U n) {
  static PrefixLagrange<T> lag;
  return lag.eval(f, n);
}

template<typename T, typename U>
T lagrange_iota_sum(const vector<T>& f, U n) {
  int m = f.size();
  vector<T> g(m + 1);
  for(int i = 1; i <= m; i++)
      g[i] = g[i-1] + f[i-1];
  return lagrange_iota(g, n);
}
} // namespace linalg
} // namespace lib

#endif

#line 1 "Lagrange.cpp"


#include <bits/stdc++.h>
#line 1 "Combinatorics.cpp"


#line 1 "BitTricks.cpp"


#line 4 "BitTricks.cpp"

namespace lib {
long long next_power_of_two(long long n) {
  if (n <= 0) return 1;
  return 1LL << (sizeof(long long) * 8 - 1 - __builtin_clzll(n) +
                 ((n & (n - 1LL)) != 0));
}
} // namespace lib


#line 5 "Combinatorics.cpp"

namespace lib {
using namespace std;
template<typename T>
struct Combinatorics {
    static vector<T> fat;
    static vector<T> inv;
    static vector<T> ifat;

    static T factorial(int i) {
        ensure_fat(next_power_of_two(i));
        return fat[i];
    }

    static T inverse(int i) {
        ensure_inv(next_power_of_two(i));
        return inv[i];
    }

    static T ifactorial(int i) {
        ensure_ifat(next_power_of_two(i));
        return ifat[i];
    }

    static T nCr(int n, int K) {
        if(K > n) return 0;
        ensure_fat(next_power_of_two(n));
        ensure_ifat(next_power_of_two(n));
        return fat[n] * ifat[n-K] * ifat[K];
    }

    static T arrangement(int n, int K) {
        return nCr(n, K) * factorial(n);
    }

    static T nCr_rep(int n, int K) {
        return interpolate(n - 1, K);
    }

    static T interpolate(int a, int b) {
        return nCr(a+b, b);
    }

    static void ensure_fat(int i) {
        int o = fat.size();
        if(i < o) return;
        fat.resize(i+1);
        for(int j = o; j <= i; j++) fat[j] = fat[j-1]*j;
    }

    static void ensure_inv(int i) {
        int o = inv.size();
        if(i < o) return;
        inv.resize(i+1);
        for(int j = o; j <= i; j++) inv[j] = -(inv[T::mod%j] * (T::mod/j));
    }

    static void ensure_ifat(int i) {
        int o = ifat.size();
        if(i < o) return;
        ifat.resize(i+1);
        ensure_inv(i);
        for(int j = o; j <= i; j++) ifat[j] = ifat[j-1]*inv[j];
    }
};

template<typename T>
vector<T> Combinatorics<T>::fat = vector<T>(1, T(1));
template<typename T>
vector<T> Combinatorics<T>::inv = vector<T>(2, T(1));
template<typename T>
vector<T> Combinatorics<T>::ifat = vector<T>(1, T(1));
} // namespace lib


#line 5 "Lagrange.cpp"

namespace lib {
using namespace std;
namespace linalg {
template <typename Field> struct PrefixLagrange {
  vector<Field> pref, suf;
  PrefixLagrange() {}

  void ensure(int n) {
    int o = pref.size();
    if (n <= o)
      return;
    pref.resize(n), suf.resize(n);
  }

  template <typename T> Field eval(const vector<Field> &v, T x) {
    using C = Combinatorics<Field>;
    assert(!v.empty());
    int d = (int)v.size() - 1;
    if (x <= d)
      return v[x];

    ensure(d + 1);

    Field a = x;
    pref[0] = suf[d] = 1;
    for (T i = 0; i < d; i++)
      pref[i + 1] = pref[i] * a, a -= 1;
    for (T i = d; i; i--)
      suf[i - 1] = suf[i] * a, a += 1;

    Field ans = 0;
    for (int i = 0; i <= d; i++) {
      Field l = pref[i] * suf[i] * C::ifactorial(i) * C::ifactorial(d-i) * v[i];
      if ((d + i) & 1)
        l = -l;
      ans += l;
    }
    return ans;
  }
};

template<typename T, typename U>
T lagrange_iota(const vector<T>& f, U n) {
  static PrefixLagrange<T> lag;
  return lag.eval(f, n);
}

template<typename T, typename U>
T lagrange_iota_sum(const vector<T>& f, U n) {
  int m = f.size();
  vector<T> g(m + 1);
  for(int i = 1; i <= m; i++)
      g[i] = g[i-1] + f[i-1];
  return lagrange_iota(g, n);
}
} // namespace linalg
} // namespace lib

Back to top page