#975146 (C++23) No.2747 Permutation Adjacent Sum

• 問題: No.2747 Permutation Adjacent Sum
• ユーザー: rogi52
• 提出日時: 2024-04-20 12:19:35
• 言語: C++23 (gcc 13.3.0 + boost 1.87.0)
• 結果: TLE
• 実行時間: -
• コード長: 21,357 bytes
• コンパイル時間: 3,209 ms
• コンパイル使用メモリ: 261,920 KB
• 実行使用メモリ: 263,140 KB
• 最終ジャッジ日時: 2024-10-12 07:38:42
• 合計ジャッジ時間: 12,848 ms
• ジャッジサーバーID （参考情報）: judge3 / judge2

ファイルパターン	結果
sample	-- * 2
other	TLE * 1 -- * 39

ソースコード

#line 2 "cp-library/src/cp-template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull  = unsigned long long;
using i32 = int;
using u32 = unsigned int;
using i64 = long long;
using u64 = unsigned long long;
using i128 = __int128_t;
template < class T > bool chmin(T& a, T b) { if(a > b) { a = b; return true; } return false; }
template < class T > bool chmax(T& a, T b) { if(a < b) { a = b; return true; } return false; }
template < class T, class U > T ceil (T x, U y) { return (x > 0 ? (x + y - 1) / y :           x / y); }
template < class T, class U > T floor(T x, U y) { return (x > 0 ?           x / y : (x - y + 1) / y); }
int popcnt(i32 x) { return __builtin_popcount(x); }
int popcnt(u32 x) { return __builtin_popcount(x); }
int popcnt(i64 x) { return __builtin_popcountll(x); }
int popcnt(u64 x) { return __builtin_popcountll(x); }

#line 2 "cp-library/src/utility/rep_itr.hpp"
template < class T > struct itr_rep {
    T i, d;
    constexpr itr_rep(const T i) noexcept : i(i), d(1) {}
    constexpr itr_rep(const T i, const T d) noexcept : i(i), d(d) {}
    void operator++() noexcept { i += d; }
    constexpr int operator*() const noexcept { return i; }
    constexpr bool operator!=(const itr_rep x) const noexcept { return d > 0 ? i < x.i : i > x.i; }
};

template < class T > struct rep {
    const itr_rep< T > s, t;
    constexpr rep(const T t) noexcept : s(0), t(t) {}
    constexpr rep(const T s, const T t) noexcept : s(s), t(t) {}
    constexpr rep(const T s, const T t, const T d) noexcept : s(s, d), t(t, d) {}
    constexpr auto begin() const noexcept { return s; }
    constexpr auto end  () const noexcept { return t; }
};

template < class T > struct revrep {
    const itr_rep < T > s, t;
    constexpr revrep(const T t) noexcept : s(t - 1, -1), t(-1, -1) {}
    constexpr revrep(const T s, const T t) noexcept : s(t - 1, -1), t(s - 1, -1) {}
    constexpr revrep(const T s, const T t, const T d) noexcept : s(t - 1, -d), t(s - 1, -d) {}
    constexpr auto begin() const noexcept { return s; }
    constexpr auto end  () const noexcept { return t; }
};
#line 3 "cp-library/src/utility/io.hpp"

/* 128bit integer */
istream& operator>>(istream& is, i128& x) {
    std::string s; is >> s;
    int pm = (s[0] == '-');
    x = 0;
    for(int i : rep(pm, int(s.size()))) x = x * 10 + (s[i] - '0');
    if(pm) x *= -1;
    return is;
}
ostream& operator<<(ostream& os, const i128& x) {
    if(x == 0) return os << '0';
    i128 y = x;
    if(y < 0) { os << '-'; y *= -1; }
    std::vector<int> ny;
    while(y > 0) { ny.push_back(y % 10); y /= 10; }
    for(int i : revrep(ny.size())) os << ny[i];
    return os;
}

template < class S, class T > istream& operator>>(istream& is,       std::pair< S, T >& x) { is >> x.first        >> x.second; return is; }
template < class S, class T > ostream& operator<<(ostream& os, const std::pair< S, T >& x) { os << x.first << " " << x.second; return os; }

namespace scanner {
    struct sca {
        template < class T > operator T() {
            T s; std::cin >> s; return s;
        }
    };
    struct vec {
        int n;
        vec(int n) : n(n) {}
        template < class T > operator std::vector< T >() {
            std::vector< T > v(n);
            for(T& x : v) std::cin >> x;
            return v;
        }
    };
    struct mat {
        int h, w;
        mat(int h, int w) : h(h), w(w) {}
        template < class T > operator std::vector< std::vector< T > >() {
            std::vector m(h, std::vector< T >(w));
            for(std::vector< T >& v : m) for(T& x : v) std::cin >> x;
            return m;
        }
    };
    struct speedup {
        speedup() {
            std::cin.tie(0);
            std::ios::sync_with_stdio(0);
        }
    } speedup_instance;
}
scanner::sca in() { return scanner::sca(); }
scanner::vec in(int n) { return scanner::vec(n); }
scanner::mat in(int h, int w) { return scanner::mat(h, w); }

namespace printer {
    void precision(int d) { std::cout << std::fixed << std::setprecision(d); }
    void flush() { std::cout.flush(); }
}

template < class T >
ostream& operator<<(ostream& os, const std::vector< T > a) {
    int n = a.size();
    for(int i : rep(n)) { os << a[i]; if(i != n - 1) os << ' '; }
    return os;
}

int print() { std::cout << '\n'; return 0; }
template < class head, class... tail > int print(head&& h, tail&&... t) {
    std::cout << h; if(sizeof...(tail)) std::cout << ' ';
    return print(std::forward<tail>(t)...);
}
template < class T > int print_n(const std::vector< T > a) {
    int n = a.size();
    for(int i : rep(n)) std::cout << a[i] << "\n";
    return 0;
}


#line 2 "cp-library/src/utility/key_val.hpp"

template < class K, class V >
struct key_val {
    K key; V val;
    key_val() {}
    key_val(K key, V val) : key(key), val(val) {}
    template < std::size_t Index >
    std::tuple_element_t< Index, key_val >& get() {
        if constexpr (Index == 0) return key;
        if constexpr (Index == 1) return val;
    }
};

namespace std {

template < class K, class V > struct tuple_size < key_val< K, V > > : integral_constant< size_t, 2 > {};
template < class K, class V > struct tuple_element < 0, key_val< K, V > > { using type = K; };
template < class K, class V > struct tuple_element < 1, key_val< K, V > > { using type = V; };

}
#line 2 "cp-library/src/utility/vec_op.hpp"
template < class T > key_val< int, T > max_of(const vector< T >& a) {
    int i = std::max_element(a.begin(), a.end()) - a.begin();
    return {i, a[i]};
}
template < class T > key_val< int, T > min_of(const vector< T >& a) {
    int i = std::min_element(a.begin(), a.end()) - a.begin();
    return {i, a[i]};
}
template < class S, class T > S sum_of(const vector< T >& a) {
    S sum = 0;
    for(const T x : a) sum += x;
    return sum;
}
template < class S, class T > vector< S > freq_of(const vector< T >& a, T L, T R) {
    vector< S > res(R - L, S(0));
    for(const T x : a) res[x - L] += 1;
    return res;
}
template < class S, class T > struct prefix_sum {
    vector< S > s;
    prefix_sum(const vector< T >& a) : s(a) {
        s.insert(s.begin(), S(0));
        for(int i : rep(a.size())) s[i + 1] += s[i];
    }
    // [L, R)
    S sum(int L, int R) { return s[R] - s[L]; }
};
#line 3 "cp-library/src/utility/heap.hpp"

template < class T > using heap_min = std::priority_queue< T, std::vector< T >, std::greater< T > >;
template < class T > using heap_max = std::priority_queue< T, std::vector< T >, std::less< T > >;

#line 27 "cp-library/src/cp-template.hpp"

#line 1 "cp-library/src/algorithm/bin_search.hpp"
template < class T, class F >
T bin_search(T ok, T ng, F f) {
    while(abs(ng - ok) > 1) {
        T mid = (ok + ng) / 2;
        (f(mid) ? ok : ng) = mid;
    }
    return ok;
}

template < class T, class F >
T bin_search_real(T ok, T ng, F f, int step = 80) {
    while(step--) {
        T mid = (ok + ng) / 2;
        (f(mid) ? ok : ng) = mid;
    }
    return ok;
}
#line 2 "cp-library/src/algorithm/argsort.hpp"

template < class T > std::vector< int > argsort(const std::vector< T > &a) {
    std::vector< int > ids((int)a.size());
    std::iota(ids.begin(), ids.end(), 0);
    std::sort(ids.begin(), ids.end(), [&](int i, int j) {
        return a[i] < a[j] || (a[i] == a[j] && i < j);
    });
    return ids;
}
#line 1 "macro.hpp"
namespace macro {

using size_type = int;
template < class container > void  sort(container& a) { std::sort(std:: begin(a), std:: end(a)); }
template < class container > void rsort(container& a) { std::sort(std::rbegin(a), std::rend(a)); }
template < class container > void reverse(container& a) { std::reverse(std::begin(a), std::end(a)); }
template < class container > void unique(container& a) {
    std::sort(std::begin(a), std::end(a));
    a.erase(std::unique(std::begin(a), std::end(a)), std::end(a));
}
template < class container > container  sorted(const container& a) { container b = a;  sort(b); return std::move(b); }
template < class container > container rsorted(const container& a) { container b = a; rsort(b); return std::move(b); }
template < class container, class compare > void sort(container& a, const compare& cmp) { std::sort(std::begin(a), std::end(a), cmp); }
template < class container, class compare > container sorted(const container& a, const compare& cmp) { container b = a; sort(b, cmp); return std::move(b); }
template < class container, class value > size_type lower_bound(const container& a, const value& x) { return std::lower_bound(std::begin(a), std::end(a), x) - std::begin(a); }
template < class container, class value > size_type upper_bound(const container& a, const value& x) { return std::upper_bound(std::begin(a), std::end(a), x) - std::begin(a); }

const std::vector<std::pair<size_type, size_type>> dir4 = { {+1,  0}, {-1,  0}, { 0, +1}, { 0, -1} };
const std::vector<std::pair<size_type, size_type>> dir8 = { {-1, -1}, {-1,  0}, {-1, +1}, { 0, -1}, { 0, +1}, {+1, -1}, {+1,  0}, {+1, +1} };

#ifdef _DEBUG
#define debug(x) std::cout << "[" << __LINE__ << "] " << #x << ": " << x << std::endl
#else
#define debug(x)
#endif

template < class container > void concat(container& a, const container& b) {
    a.insert(std::end(a), std::begin(b), std::end(b));
}
std::vector<size_type> iota(const size_type n) {
    std::vector<size_type> I(n);
    std::iota(std::begin(I), std::end(I), 0);
    return I;
}
template < class container > std::vector<size_type> sort_idx(const container& a) {
    const size_type n = a.size();
    std::vector<size_type> I = iota(n);
    std::sort(std::begin(I), std::end(I), [&](size_type i, size_type j) { return a[i] < a[j] or (a[i] == a[j] and i < j); });
    return I;
}
template < class container, class compare > std::vector<size_type> sort_idx(const container& a, const compare& cmp) {
    const size_type n = a.size();
    std::vector<size_type> I = iota(n);
    std::sort(std::begin(I), std::end(I), [&](size_type i, size_type j) { return cmp(a[i], a[j]) or (a[i] == a[j] and i < j); });
    return std::move(I);
}

struct grid {
    using size_type = int;
    size_type H, W;
    grid(const size_type H, const size_type W) : H(H), W(W) {}
    bool contains(const size_type i, const size_type j) {
        return 0 <= i and i < H and 0 <= j and j < W;
    }
};

using f64 = long double;

template < class T > vector< T >& operator++(vector< T >& a) { for(T& x : a) x++; return a; }
template < class T > vector< T >& operator--(vector< T >& a) { for(T& x : a) x--; return a; }
template < class T > vector< T >  operator++(vector< T >& a, signed) { vector< T > res = a; for(T& x : a) x++; return res; }
template < class T > vector< T >  operator--(vector< T >& a, signed) { vector< T > res = a; for(T& x : a) x--; return res; }

} // namespace macro

using namespace macro;
#line 2 "cp-library/src/number/modint.hpp"
struct modinfo { uint mod, root, isprime; };
template < modinfo const &ref >
struct modint {
    static constexpr uint const &mod = ref.mod;
    static constexpr uint const &root = ref.root;
    static constexpr uint const &isprime = ref.isprime;
    uint v = 0;
    constexpr modint& s(uint v) { this->v = v < mod ? v : v - mod; return *this; }
    constexpr modint(ll v = 0) { s(v % mod + mod); }
    modint operator-() const { return modint() - *this; }
    modint& operator+=(const modint& rhs) { return s(v + rhs.v); }
    modint& operator-=(const modint& rhs) { return s(v + mod - rhs.v); }
    modint& operator*=(const modint& rhs) { v = ull(v) * rhs.v % mod; return *this; }
    modint& operator/=(const modint& rhs) { return *this *= inv(rhs); }
    modint operator+(const modint& rhs) const { return modint(*this) += rhs; }
    modint operator-(const modint& rhs) const { return modint(*this) -= rhs; }
    modint operator*(const modint& rhs) const { return modint(*this) *= rhs; }
    modint operator/(const modint& rhs) const { return modint(*this) /= rhs; }
    friend modint pow(modint x, ll n) { modint res(1); while(n > 0) { if(n & 1) res *= x; x *= x; n >>= 1; } return res; }
    friend modint inv(modint v) {
        if(isprime) {
            return pow(v, mod - 2);
        } else {
            ll a = v.v, b = modint::mod, x = 1, y = 0, t;
            while(b > 0) { t = a / b; swap(a -= t * b, b); swap(x -= t * y, y); }
            return modint(x);
        }
    }
    friend modint operator+(int x, const modint& y) { return modint(x) + y; }
    friend modint operator-(int x, const modint& y) { return modint(x) - y; }
    friend modint operator*(int x, const modint& y) { return modint(x) * y; }
    friend modint operator/(int x, const modint& y) { return modint(x) / y; }
    friend istream& operator>>(istream& is, modint& m) { ll x; is >> x; m = modint(x); return is; }
    friend ostream& operator<<(ostream& os, const modint& m) { return os << m.v; }
    bool operator==(const modint& r) const { return v == r.v; }
    bool operator!=(const modint& r) const { return v != r.v; }
    static uint get_mod() { return mod; }
    static int is_prime() { return isprime; }
};
constexpr modinfo base998244353 { 998244353, 3, 1 };
constexpr modinfo base1000000007 { 1000000007, 0, 1 };
using mint998244353 = modint< base998244353 >;
using mint1000000007 = modint< base1000000007 >;
#line 3 "cp-library/src/number/binom_mod.hpp"

template < class mint >
mint fact(int n) {
    assert(0 <= n);
    assert(mint::is_prime());
    static const uint mod = mint::get_mod();
    static std::vector<mint> data = {1, 1};
    while(int(data.size()) <= n) {
        int i = data.size();
        data.push_back(data.back() * i);
    }
    return data[n];
}

template < class mint >
mint inv(int n) {
    assert(0 <= n);
    assert(mint::is_prime());
    static const uint mod = mint::get_mod();
    static std::vector<mint> data = {1, 1};
    while(int(data.size()) <= n) {
        int i = data.size();
        data.push_back(- data[mod % i] * (mod / i));
    }
    return data[n];
}

template < class mint >
mint fact_inv(int n) {
    assert(0 <= n);
    assert(mint::is_prime());
    static const uint mod = mint::get_mod();
    static std::vector<mint> data = {1, 1};
    while(int(data.size()) <= n) {
        int i = data.size();
        data.push_back(data.back() * inv<mint>(i));
    }
    return data[n];
}

template < class mint >
mint comb(int n, int k) {
    if(k < 0 or n < k) return 0;
    return fact<mint>(n) * fact_inv<mint>(k) * fact_inv<mint>(n - k);
}

template < class mint >
mint perm(int n, int k) {
    return fact<mint>(n) * fact_inv<mint>(n - k);
}

template < class mint >
mint homo(int n, int k) {
    return comb<mint>(n + k - 1, k);
}

template < class mint > struct power {
    mint a;
    std::vector<mint> data = {1};
    power() {}
    power(const mint a) : a(a) {}
    // a^n
    mint get(const int n) {
        assert(0 <= n);
        while(int(data.size()) <= n)
            data.push_back(data.back() * a);
        return data[n];
    }
};
#line 5 "C.cpp"

#line 7 "C.cpp"
#define rep(i,n) for(int i = 0; i < (n); i++)
using namespace std;
using ll = long long;
using ld = long double;

template< int mod >
struct Fp {
    ll x;
    Fp() : x(0) {}
    Fp(int64_t y) : x(y >= 0 ? y % mod : (mod + y % mod) % mod) {}
    Fp &operator+=(const Fp &p) { if((x += p.x) >= mod) x -= mod; return *this; }
    Fp &operator-=(const Fp &p) { if((x += mod - p.x) >= mod) x -= mod; return *this; }
    Fp &operator*=(const Fp &p) { x = (int)(1LL * x * p.x % mod); return *this; }
    Fp &operator/=(const Fp &p) { *this *= p.inv(); return *this; }
    Fp operator-() const { return Fp(-x); }
    Fp operator+(const Fp &p) const { return Fp(*this) += p; }
    Fp operator-(const Fp &p) const { return Fp(*this) -= p; }
    Fp operator*(const Fp &p) const { return Fp(*this) *= p; }
    Fp operator/(const Fp &p) const { return Fp(*this) /= p; }
    bool operator==(const Fp &p) const { return x == p.x; }
    bool operator!=(const Fp &p) const { return x != p.x; }

    Fp inv() const {
        int a = x, b = mod, u = 1, v = 0, t;
        while(b > 0) {
            t = a / b;
            swap(a -= t * b, b);
            swap(u -= t * v, v);
        }
        return Fp(u);
    }

    Fp pow(int64_t n) const {
        Fp ret(1), mul(x);
        while(n > 0) {
            if(n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }

    friend ostream &operator<<(ostream &os, const Fp &p) { return os << p.x; }
    friend istream &operator>>(istream &is, Fp &a) {
        int64_t t;
        is >> t;
        a = Fp< mod > (t);
        return (is);
    }

    static long long getmod() { return mod; }
};

using mint = Fp<998244353>;

// Binomial coefficient
template<class T> struct BiCoef {
    vector<T> fact_, inv_, finv_;
    constexpr BiCoef() {}
    constexpr BiCoef(int n) noexcept : fact_(n, 1), inv_(n, 1), finv_(n, 1) {
        init(n);
    }
    constexpr void init(int n) noexcept {
        fact_.assign(n, 1), inv_.assign(n, 1), finv_.assign(n, 1);
        int MOD = fact_[0].getmod();
        for(int i = 2; i < n; i++){
            fact_[i] = fact_[i-1] * i;
            inv_[i] = -inv_[MOD%i] * (MOD/i);
            finv_[i] = finv_[i-1] * inv_[i];
        }
    }
    constexpr T com(int n, int k) const noexcept {
        if (n < k || n < 0 || k < 0) return 0;
        return fact_[n] * finv_[k] * finv_[n-k];
    }
    constexpr T fact(int n) const noexcept {
        if (n < 0) return 0;
        return fact_[n];
    }
    constexpr T inv(int n) const noexcept {
        if (n < 0) return 0;
        return inv_[n];
    }
    constexpr T finv(int n) const noexcept {
        if (n < 0) return 0;
        return finv_[n];
    }
};

const int MOD = 998244353;
using mint = Fp<MOD>;

template < class S >
S lagrange_interpolation(vector< S > x, vector< S > y, S T) {
    int n = x.size() - 1u;
    S res(0);
    for(int i = 0; i <= n; i++) {
        S ft(1), fx(1);
        for(int j = 0; j <= n; j++) if(i != j) {
            ft *= T - x[j];
            fx *= x[i] - x[j];
        }
        res += y[i] * ft / fx;
    }
    return res;
}

template < class S >
S lagrange_interpolation_arithmetic(S a, S d, vector< S > y, S T) {
    int n = y.size() - 1u;
    vector< S > fx(n + 1);
    fx[0] = 1;
    for(int j = 1; j <= n; j++) fx[0] *= mint(-j) * d;
    for(int i = 0; i <= n - 1; i++) fx[i + 1] = fx[i] * (i + 1) / (i - n);

    vector< S > ft(n + 1), L(n + 1), R(n + 1);
    L[0] = R[n] = 1;
    for(int i = 0; i < n; i++) L[i + 1] = L[i] * (T - (a + S(i) * d));
    for(int i = n; i > 0; i--) R[i - 1] = R[i] * (T - (a + S(i) * d));
    for(int i = 0; i <= n; i++) ft[i] = L[i] * R[i];

    S res(0);
    for(int i = 0; i <= n; i++) {
        res += y[i] * ft[i] / fx[i];
    }
    return res;
}

template < class S >
vector< S > lagrange_interpolation(vector< S > x, vector< S > y) {
    int n = x.size() - 1u;
    vector< S > c(n + 1);
    for(int i = 0; i <= n; i++) {
        S fx(1);
        for(int j = 0; j <= n; j++) if(i != j) fx *= x[i] - x[j];
        c[i] = y[i] / fx;
    }
    // g = prod (x - x[i])
    vector< S > g(n + 2);
    g[0] = 1;
    for(int i = 0; i <= n; i++) {
        vector< S > g2(n + 2, 0), a = {-x[i], 1};
        for(int j = 0; j <= n; j++) {
            for(int k = 0; k <= 1; k++) {
                g2[j + k] += g[j] * a[k];
            }
        }
        swap(g, g2);
    }

    vector< S > xinv(n + 1);
    for(int i = 0; i <= n; i++) xinv[i] = S(1) / x[i];
    vector< S > f(n + 1);
    for(int i = 0; i <= n; i++) {
        vector< S > g2 = g, h(n + 1), a = {-x[i], 1};
        for(int j = n; j >= 0; j--) {
            h[j] = g2[j + 1];
            f[j] += h[j] * c[i];
            for(int k = 0; k <= 1; k++) g2[k + j] -= a[k] * h[j];
        }
    }
    return f;
}


BiCoef<mint> bc(10001000);

mint lagrange_interpolation_0n(vector< mint > y, mint T) {
    int n = y.size() - 1u;
    if(0 <= T.x && T.x <= n) return y[T.x];

/*
    vector< mint > fact(n + 1, 1), inv(n + 1, 1), finv(n + 1, 1);
    int MOD = fact[0].getmod();
    for(int i = 2; i <= n; i++) {
        fact[i] = fact[i - 1] * i;
        inv[i] = -inv[MOD % i] * (MOD / i);
        finv[i] = finv[i - 1] * inv[i];
    }
*/

    vector< mint > L(n + 1), R(n + 1);
    L[0] = R[n] = 1;
    for(int i = 0; i < n; i++) L[i + 1] = L[i] * mint(T - i);
    for(int i = n; i > 0; i--) R[i - 1] = R[i] * mint(T - i);
    mint res(0);
    for(int i = 0; i <= n; i++) {
        res += y[i] * L[i] * R[i] * bc.finv(i) * bc.finv(n - i) * ((n - i) % 2 ? -1 : +1);
    }
    return res;
}

mint solve(i64 r, i64 d, i64 n) {
    if(n == 0) { return 0; }
    n--;

    vector<mint> sum(d + 2);
    rep(i,d+2) sum[i] = mint(r).pow(i) * mint(i).pow(d);
    rep(i,d+1) sum[i + 1] += sum[i];
    if(r == 1) {
        return lagrange_interpolation_0n(sum, n);
    } else {
        mint c = 0;
        rep(i,d+1) c += bc.com(d + 1, i + 1) * mint(-r).pow(d - i) * sum[i];
        c /= mint(1 - r).pow(d + 1);
        vector<mint> y(d + 1);
        mint div = 1, rinv = mint(r).inv();
        rep(i,d+1) { y[i] = (sum[i] - c) * div; div *= rinv; }
        return mint(r).pow(n) * lagrange_interpolation_0n(y, n) + c;
    }
}

int main() {
    int N = in(), K = in();
    mint ans = 0;
    ans += solve(1, K, N) * N - solve(1, K + 1, N);
    print(ans * bc.fact(N - 1) * 2);
}