[백준] 16978번 수열과 쿼리 22 C++ 코드

코드

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef long long ll;

const int MAX_N = 100003;
ll n, q, m, cnt = 0;
ll a[MAX_N];
vector<ll> tree(MAX_N + 1, 0);
vector<pair<pair<int, int>, pair<int, int>>> queries;
vector<pair<int, int>> updates;
vector<pair<int, ll>> res;

void update(int idx, ll diff) {
    while (idx <= n) {  // 범위를 n으로 제한
        tree[idx] += diff;
        idx += (idx & -idx);
    }
}

ll sum(int idx) {
    ll ans = 0;
    while (idx > 0) {
        ans += tree[idx];
        idx -= (idx & -idx);
    }
    return ans;
}

void FASTIO() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
}

bool cmp(pair<pair<int, int>, pair<int, int>> a, pair<pair<int, int>, pair<int, int>> b) {
    return a.first.first < b.first.first;
}

int main() {
    FASTIO();
    cin >> n;

    // 배열 초기화 및 Fenwick Tree 초기화
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
        update(i, a[i]);
    }

    cin >> m;
    for (int e = 0; e < m; e++) {
        int q;
        cin >> q;
        if (q == 1) {
            int i, v;
            cin >> i >> v;
            updates.push_back({i, v});
        } else if (q == 2) {
            int k, i, j;
            cin >> k >> i >> j;
            queries.push_back({{k, e}, {i, j}});
        }
    }

    sort(queries.begin(), queries.end(), cmp);  // k 기준으로 정렬

    int i = 0;
    int len = queries.size();
    while (i < len) {
        if (cnt == queries[i].first.first) {  // k번째 업데이트 처리 후 결과 계산
            ll tmp = sum(queries[i].second.second) - sum(queries[i].second.first - 1);
            res.push_back({queries[i].first.second, tmp});
            i++;
        } else if (cnt < queries[i].first.first) {  // 업데이트 적용
            int idx = updates[cnt].first;
            int value = updates[cnt].second;

            update(idx, -a[idx]);  // 이전 값 제거
            update(idx, value);   // 새로운 값 추가
            a[idx] = value;       // 배열 업데이트
            cnt++;
        }
    }

    // 결과 정렬 및 출력
    sort(res.begin(), res.end());
    for (const auto& r : res) {
        cout << r.second << '\n';
    }

    return 0;
}

풀이

구간합을 구하는 문제라서 펜윅트리를 이용했다. n번째 업데이트를 적용했을 때 합을 구하라는 걸 보고 처음에 어떻게 하지 하고 포기했었는데, 오프라인 쿼리 방식을 학습한 후에 보니까 방법이 떠올랐다. vector에 구할 쿼리를 저장한다. 그 후 k를 기준으로 정렬한다. k값이 현재 cnt와 같으면 값을 구하고, k 값이 커지면 업데이트를 하고 값을 찾아주면 된다.