区間和の和

2 secs 1024 MB

shinnshinn

問題

提出

テストケース

解説

Tester解みどりむし：セグ木を使った解法

解説

考えやすくするためクエリ２の区間を半開区間 $[L,R)$ としておきます。

$\sum_{j = L}^{i}$ の部分を $S_i$ とします。

$S_i$ はそれぞれ以下のようになります。

\begin{aligned} S_L &= A_L \\ S_{L+1} &= A_L + A_{L+1}\\ \vdots \\ S_{R-1} &= A_L + A_{L+1} + \cdots + A_{R-1} \end{aligned}

よって、

\begin{aligned} \sum_{i = L}^R S_i &= \sum_{i=L}^R(R-i) \times A_i \\ &= R\sum_{i=L}^{R}A_i - \sum_{i=L}^{R}iA_i \end{aligned}

となります。それそれの総和はFenwick Treeで管理すれば総和の取得、要素の更新が高速に行えます。

実装

xxxxxxxxxx
 
#include <bits/stdc++.h>
#include <atcoder/modint>
using namespace atcoder;
using namespace std;
using mint = modint998244353;
​
template<typename T>
struct FenwickTree{
    int n;
    vector<T> data;
​
    FenwickTree(int n):n(n), data(n+1){}
​
    void add(int i, T x){
        assert(0 <= i && i < n);
        i++;
        while(i <= n){
            data[i] += x;
            i += i & -i;
        }
    }
​
    T get(int i) const {
        assert(0 <= i && i < n);
        return sum(i, i+1);
    }
​
    // sum[0, i)
    T sum(int i) const{
        assert(0 <= i && i <= n);
        T retu_sum = 0;
        while(i > 0){
            retu_sum += data[i];
            i -= i & -i;
        }
​
        return retu_sum;
    }
​
    // sum[l, r)
    T sum(int l, int r) const {
        assert(0 <= l && l <= n);
        assert(0 <= r && r <= n);
        return sum(r) - sum(l);
    }
};
​
void solve() {
    int n;
    cin >> n;
    vector<mint> a(n);
    for(int i = 0;i < n; ++i){
        int val;
        cin >> val;
        a[i] = val;
    }
​
    FenwickTree<mint> tree1(n+1), tree2(n+1);
    for(int i = 0;i < n; ++i){
        tree1.add(i, a[i]);
        tree2.add(i, a[i] * i);
    }
​
    int q;
    cin >> q;
    while(q--){
        int query;
        cin >> query;
        if(query == 1){
            int l, r;
            cin >> l >> r;
            --l;
            
            mint ans = r*tree1.sum(l, r) - tree2.sum(l, r);
            cout << ans.val() << '\n';
        }else{
            int k, x;
            cin >> k >> x;
            --k;
            tree1.add(k, mint(x) - tree1.get(k));
            tree2.add(k, mint(x)*k - tree2.get(k));
​
        }
    }
} 
​
int main() {
    ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    solve();
    return 0;
}
​