BOJ 16978 수열과 쿼리 22

29분 26초
세그먼트 트리 문제인데 query의 순서를 정렬하면 해결된다.
합을 출력할 때 k의 순서가 오름차순이 되도록 정렬하고 ans 배열을 만들어 원래의 인덱스에 값을 넣어 원래 순서대로 출력하면 된다.

조심해야 할 부분은 구간합이 10^5 * 10^6 == 10^11으로 int 범위가 넘어가서 long long을 사용해야 한다. ans 배열도 long long으로 해서 출력해야 한다.

#include <iostream>
#include <algorithm>
#include <cmath>
#include <utility>
#include <string>
#include <cstring>
#include <vector>
#include <tuple>
#include <stack>
#include <queue>
#include <deque>
#include <list>
#include <map>
#include <unordered_map>
#include <climits>

#define INF 987654321
#define INF2 2147483647
#define f first
#define s second
#define x first
#define y second
#define all(v) (v).begin(), (v).end()

using namespace std;
using ll = long long;
using pii = pair<int, int>;

class segment {
public:
    vector<ll> tree; //tree[node] := a[start ~ end] 의 합

    segment() {}
    segment(int size) {
        this->resize(size);
    }
    void resize(int size) {
        size = (int) floor(log2(size)) + 2;
        size = pow(2, size);
        tree.resize(size);
    }
    ll init(vector<ll> &a, int node, int start, int end) {
        if(start == end) return tree[node] = a[start];
        else return tree[node] = init(a, 2 * node, start, (start + end) / 2) +
                                 init(a, 2 * node + 1, (start + end) / 2 + 1, end);
    }
    ll sum(int node, int start, int end, int left, int right) {
        if(right < start || end < left) return 0;
        if(left <= start && end <= right) return tree[node];
        return sum(node * 2, start, (start + end) / 2, left, right) +
               sum(node * 2 + 1, (start + end) / 2 + 1, end, left, right);
    }
    void update(int node, int start, int end, int index, ll diff) {
        if(index < start || end < index) return;
        tree[node] += diff;
        if(start != end) {
            update(node * 2, start, (start + end) / 2, index, diff);
            update(node * 2 + 1, (start + end) / 2 + 1, end, index, diff);
        }
    }
};

typedef struct type1 {
    int i; ll v;
    type1() {}
    type1(int pm1, int pm2) : i(pm1), v(pm2) {}
};
typedef struct type2 {
    int k, i, j;
    int idx;
    type2(int pm1, int pm2, int pm3, int pm4) : k(pm1), i(pm2), j(pm3), idx(pm4) {}
};
bool cmp(type2 a, type2 b) {
    return a.k  < b.k;
}

int N, M;
vector<ll> A;
vector<type1> query1(1);
vector<type2> query2;
int query2Size = 0;
vector<ll> ans;

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    cin >> N;
    A.resize(N+1);
    segment root(N);
    for(int i=1; i<=N; i++) cin >> A[i];
    root.init(A, 1, 1, N);
    cin >> M;
    while(M--) {
        int q; cin >> q;
        if(q == 1) {
            int i, v; cin >> i >> v;
            query1.emplace_back(i,v);
        }
        else {
            int k, i, j; cin >> k >> i >> j;
            query2.emplace_back(k,i,j,query2Size++);
        }
    }
    ans.resize(query2Size);
    sort(all(query2), cmp);
    int q1 = 1;
    for(type2 q2 : query2) {
        while(q1 <= q2.k) {
            root.update(1,1,N,query1[q1].i, query1[q1].v - A[query1[q1].i]);
            A[query1[q1].i] = query1[q1].v;
            q1++;
        }
        ans[q2.idx] = root.sum(1,1,N,q2.i,q2.j);
    }

    for(ll x : ans) {
        cout << x << '\n';
    }
    return 0;
}