[BOJ / C++ / Pyton3] 2042번 : 구간합 구하기

 

문제 링크

https://www.acmicpc.net/problem/2042

 

2 초

256 MB

120606

32605

16662

26.465%

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

 

Solved.ac 레벨

골드 I

문제 풀이 해설

python 풀이 아래 포스팅에서 자세히 설명해두었습니다.

C++의 경우 주석 참고.

https://heibondk.tistory.com/537

[4주차] 세그먼트 트리 (Segment tree)

 

문제 풀이 코드

#include <bits/stdc++.h>
using namespace std;

typedef long long int ll;

// n : 수의 개수
// m : 수의 변경이 일어나는 횟수
// k : 구간의 합을 구하는 횟수
int n, m, k;
vector<ll> arr;
vector<ll> segTree;

ll init(vector<ll> &arr, vector<ll> &segTree, int node, int start, int end);
void update(vector<ll> &segTree, int node, int start, int end, int idx, ll diff);
ll prefix(vector<ll> &segTree, int node, int start, int end, int left, int right);

int main(){
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    cin >> n >> m >> k;

    arr.resize(n + 1);
    segTree.resize(4 * n);

    for(int i = 1; i <= n; i ++) {
        cin >> arr[i];
    }

    // 세그먼트 트리 초기화
    init(arr, segTree, 1, 1, n);

    for(int i = 0; i < m+k; i++) {
        ll a, b, c;
        cin >> a >> b >> c;

        if(a == 1) {
            // a가 1이면 갱신
            ll diff = c - arr[b];
            arr[b] = c;
            update(segTree, 1, 1, n, b, diff);
        } else if (a == 2) {
            // a가 2이면 구간합 
            cout << prefix(segTree, 1, 1, n, b, c) << "\n";
        }
    }

    return 0;
}

ll init(vector<ll> &arr, vector<ll> &segTree, int node, int start, int end) {
    // 노드가 리프노드인 경우
    if(start == end) 
        return segTree[node] = arr[start];

    // 현재 노드가 담당하는 구간을 정확히 반으로 나눔
    int mid = (start + end) / 2;

    return segTree[node] = init(arr, segTree, node * 2, start, mid) + init(arr, segTree, node * 2 + 1, mid + 1, end);
}

ll prefix(vector<ll> &segTree, int node, int start, int end, int left, int right) {
    // [start, end] 앞 뒤에 [left, right]가 있는 경우
    // 겹치지 않으므로 탐색을 더 이상 할 필요가 없음.
    if(left > end || right < start) return 0;

    // [start, end]가 [left, right]에 포함되는 경우
    if(left <= start && end <= right) return segTree[node];

    // 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서 재탐색
    int mid = (start + end) / 2;

    return prefix(segTree, node * 2, start, mid, left, right) + prefix(segTree, node * 2 + 1, mid + 1, end, left, right);
}

void update(vector<ll> &segTree, int node, int start, int end, int idx, ll diff) {
    if (idx < start || idx > end) return;

    segTree[node] = segTree[node] + diff;

    // 리프 노드가 아닌 경우 자식도 변경해줘야함
    // 리프 노드인지 검사를 하고 아래 자식 노드 갱신
    // start == end => 리프 노드
    if(start != end) {
        int mid = (start + end) / 2;
        update(segTree, node * 2, start, mid, idx, diff);
        update(segTree, node * 2 + 1, mid + 1, end, idx, diff);
    }
}