F - go up stairs (hard)

2 secs 1024 MB
kusirakusira's icon kusirakusira

ここで
M=max(A)M = max(A)
DPi:=iDP_i := i段目に行く通りの数
と定義します。
X>MX > Mである時のDPXM,DPXM+1,,DPX1DP_{X-M},DP_{X-M+1},\cdots ,DP_{X-1}をE_2よりも高速に求めることを考えます。
ここでkitamasa法を用いることでDPXMDP_{X-M}O(M2logX)O(M^2logX)で求めることができます。
また、kitamasa法ではDPnDP_{n}の結果からDPn+1DP_{n+1}O(M)O(M)で求めることができるので DPXKDP_{X-K}の結果からDPXM+1,DPXM+2,,DPX1DP_{X-M+1},DP_{X-M+2},\cdots ,DP_{X-1}O(M2)O(M^2)で求めることができます。
このようにしてこの問題をO(M2logX)O(M^2logX)で解くことができます。
また、MXM \leq Xの場合 である場合もDP0,,DPX1DP_0,\cdots ,DP_{X-1}を同様の方法で求めることができます。
よすぽさんのkitamasa法に関する記事

python3
import sys

def main():
    input = sys.stdin.buffer.readline

    def plus_one(c0, c, k=1000, mod=998244353):
        next = [0] * k
        next[0] = c[k - 1] * c0[0] % mod
        for i in range(1, k):
            next[i] = (c[i - 1] + c[k - 1] * c0[i]) % mod

        for i in range(k):
            c[i] = next[i]

    def double(c0, c, k=1000, mod=998244353):
        mat = [0] * (k * k)
        now = [i for i in c]
        for i in range(k):
            for j in range(k):
                mat[i * k + j] = now[j]
            if i < k - 1:
                plus_one(c0, now, k, mod)

        next = [0] * k
        for m in range(k * k):
            i = m // k
            j = m % k
            next[i] += mat[j] * mat[j * k + i] % mod

        for i in range(k):
            c[i] = next[i]

    def kitamasa(c0, n, k=1000, mod=998244353):
        actions = []
        n += k - 1
        while n >= 2 * k:
            if n & 1 == 1:
                actions.append(0)
                n ^= 1
            else:
                actions.append(1)
                n >>= 1

        while n > k:
            actions.append(0)
            n -= 1

        actions.reverse()
        c = [i for i in c0]

        for i in actions:
            if i == 0:
                plus_one(c0, c, k, mod)
            else:
                double(c0, c, k, mod)

        return c

    n, x = map(int, input().split())
    a = list(map(int, input().split()))
    k = 300
    mod = 998244353

    c0 = [0] * k
    for i in range(n):
        c0[k - a[i]] += 1

    dp = [0] * k
    c = []
    is_first = True
    for i in range(k):
        if x - k + i < 0:
            continue
        if x - k + i == 0:
            dp[i] = 1
            continue
        if is_first:
            c = kitamasa(c0, x - k + i, k, mod)
            is_first = False
        else:
            plus_one(c0, c, k, mod)

        dp[i] = c[-1]

    ans = 0

    for m in range(k * n):
        i = m // n
        j = m % n
        if i + a[j] >= k:
            ans += dp[i]

    print(ans % mod)


if __name__ == "__main__":
    main()
C++
#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;

using ll = long long;
using mint = modint998244353;  // 1000000007;

// O(K)
void plus_one(int k, const vector<mint>& c0, vector<mint>& c) {
    vector<mint> next(k);
    next[0] = c[k - 1] * c0[0];
    for (int i = 1; i < k; i++) {
        next[i] = c[i - 1] + c[k - 1] * c0[i];
    }
    c.swap(next);
}

// O(K^2)
void make_double(int k, const vector<mint>& c0, vector<mint>& c) {
    vector<vector<mint>> cs(k);
    cs[0] = c;
    for (int i = 1; i < k; i++) {
        cs[i] = cs[i - 1];
        plus_one(k, c0, cs[i]);
    }
    vector<mint> next(k);
    for (int i = 0; i < k; i++) {
        for (int j = 0; j < k; j++) {
            next[i] += cs[0][j] * cs[j][i];
        }
    }
    c.swap(next);
}

void kitamasa(int k, const vector<mint> c0, vector<mint>& c, ll n) {
    vector<bool> action;
    while (n >= k * 2) {
        if (n & 1) {
            action.emplace_back(0);
            n -= 1;
        } else {
            action.emplace_back(1);
            n >>= 1;
        }
    }
    while (n > k) {
        n--;
        action.emplace_back(0);
    }

    reverse(action.begin(), action.end());
    for (int i = 0; i < (int)action.size(); i++) {
        if (action[i]) {
            make_double(k, c0, c);
        } else {
            plus_one(k, c0, c);
        }
    }
}

int main() {
    ll n, x;
    int k = 300;
    cin >> n >> x;
    assert(n <= k);
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        assert(a[i] <= k);
    }
    vector<mint> c(k, 0);
    for (int i = 0; i < n; i++) c[k - a[i]]++;
    const auto c0 = c;

    // ふつうのkitamasa法を用いてるためO(K^2 log n)
    // もし高速kitamasa法を用いた場合はO(K^2 + K logk logn)

    vector<mint> bef(k, 0);
    if (x <= k) {
        bef[x - 1] = 1;
        for (int i = 0; i < x - 1; i++) {
            bef[x - 2 - i] = c[k - 1];
            plus_one(k, c0, c);
        }
    } else {
        kitamasa(k, c0, c, x - 1);
        for (int i = 0; i < k; i++) {
            bef[k - 1 - i] = c[k - 1];
            plus_one(k, c0, c);
        }
    }

    mint ans = 0;
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < a[i]; j++) {
            ans += bef[j];
        }
    }
    cout << ans.val() << endl;
}

別解

Bostan-Mori法でもAC出来ます。