Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code: Sum of Fractions

#include <bits/stdc++.h>

constexpr int P = 998244353;

typedef long long ll;

int power(int x, int y = P - 2) {
    int ans = 1;
    while (y) {
        if (y & 1) {
            ans = (ll)ans * x % P;
        }
        x = (ll)x * x % P;
        y >>= 1;
    }
    return ans;
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    int n, q;
    std::cin >> n >> q;
    std::vector<int> a(n), inv(n);
    std::vector<int> sinv(n + 1), sinvi(n + 1);
    int k = std::__lg(n);
    std::vector st(k + 1, std::vector<std::pair<int, int>>(n));
    ll s1 = 0, s2 = 0;
    for (int i = 0; i < n; i++) {
        std::cin >> a[i];
        inv[i] = power(a[i]);
        st[0][i] = std::make_pair(a[i], i);
        s1 += inv[i], s2 += (ll)inv[i] * i % P;
        sinv[i + 1] = s1 % P, sinvi[i + 1] = s2 % P;
    }
    for (int i = 1; i <= k; i++) {
        for (int j = 0; j + (1 << i) - 1 <= n; j++) {
            st[i][j] = std::min(st[i - 1][j], st[i - 1][j + (1 << (i - 1))]);
        }
    }
    auto qmin = [&](int l, int r) -> int {
        int k = std::__lg(r - l + 1);
        return std::min(st[k][l], st[k][r - (1 << k) + 1]).second;
    };
    std::vector<int> mul(n);
    ll base = 0;
    auto dfs = [&](auto &&self, int l, int r) -> void {
        int p = qmin(l, r);
        mul[p] = (ll)(p - l + 1) * (r - p + 1) % P;
        ll s1 = 0, s2 = 0;
        s1 += sinvi[p] - sinvi[l] + P;
        s1 += (ll)(P - l + 1) * (sinv[p] - sinv[l] + P);
        s2 += sinvi[p + 1] - sinvi[r + 1] + P;
        s2 += (ll)(r + 1) * (sinv[r + 1] - sinv[p + 1] + P);
        // for (int i = l; i < p; i++) s1 += (ll)(i - l + 1) * inv[i] % P;
        // for (int i = p + 1; i <= r; i++) s2 += (ll)(r - i + 1) * inv[i] % P;
        base += s1 % P * (r - p + 1) % P;
        base += s2 % P * (p - l + 1) % P;
        if (l < p)
            self(self, l, p - 1);
        if (p < r)
            self(self, p + 1, r);
    };
    dfs(dfs, 0, n - 1);
    base %= P;
    std::vector<std::tuple<int, int, int>> vec(n);
    for (int i = 0; i < n; i++) {
        vec[i] = std::make_tuple(a[i], inv[i], mul[i]);
    }
    std::sort(vec.begin(), vec.end());
    std::vector<int> f(n + 1), g1(n + 1), g2(n + 1);
    for (int i = n - 1; i >= 0; i--) {
        auto [a, inv, mul] = vec[i];
        f[i] = (f[i + 1] + (ll)inv * mul) % P;
    }
    for (int i = 0; i < n; i++) {
        auto [a, inv, mul] = vec[i];
        g1[i + 1] = (g1[i] + (ll)mul * a) % P;
        g2[i + 1] = (g2[i] + mul) % P;
    }
    while (q--) {
        int x;
        std::cin >> x;
        ll ans = base;
        int p = std::lower_bound(vec.begin(), vec.end(),
                                 std::make_tuple(x + 2, 0, 0)) -
                vec.begin();
        ans += (ll)f[p] * (x + 1);
        ans += (ll)g2[p] * (x + 2);
        ans += P - g1[p];
        // for (int i = 0; i < n; i++) {
        //     int val;
        //     if (x + 2 <= a[i]) {
        //         // val = (ll)(1 + x) * inv[i] * mul[i];
        //         val = 0;
        //     }
        //     else {
        //         val = 2 + x - a[i];
        //     }
        //     ans += (ll)val * mul[i] % P;
        // }
        std::cout << (ans % P) << std::endl;
    }
    return 0;
}