[BOJ 2042] 구간 합 구하기

2042번: 구간 합 구하기

1. 문제 설명

 

이번 문제는 일반적인 구간 합 문제와 달리 수의 변경이 빈번하게 일어나는 것이 핵심적인 차이점이다.

간단한 구간 합 문제는 누적 합의 차로 쉽게 구할 수 있지만, 중간에 계속 수가 변경된다면 누적 합이 일정하지 않으므로 해당 방법을 사용할 수 없다.

 

이렇게 원소가 변경될 때 구간 합을 구할 때 사용할 수 있는 방법이 세그먼트 트리이다. 

자세한 개념 설명은 따로 정리해두었으니 아래 링크를 참고하자.

 

[C++] 세그먼트 트리의 개념과 구현

 

2. 구현 아이디어 1 - 틀렸습니다

세그먼트 트리 알고리즘을 그대로 적용해주면 되는 문제라고 생각하고 문제를 풀었다.

여기서 index를 의미하는 b가 1부터 시작한다는 것을 주의하여 입력을 받은 후 b를 1 줄여주었다.

b는 1부터 시작한다

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

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

    int mid = (start + end) / 2;

    ll leftChild = init(arr, tree, node*2, start, mid);
    ll rightChild = init(arr, tree, node*2+1, mid + 1, end);

    // 구간 합을 구하는 경우
    return tree[node] = leftChild + rightChild;
}

ll query(vector<ll> &tree, int node, int start, int end, int left, int right) {
    // case 1: [start, end] 앞 뒤에 [left, right]가 있는 경우
    if (left > end || right < start) return 0;

    // case 2: [start, end]가 [left, right]에 포함
    if (left <= start && end <= right) return tree[node];

    // case 3, 4: 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서 다시 탐색 시작
    int mid = (start + end) / 2;
    ll leftChild = query(tree, node*2, start, mid, left, right);
    ll rightChild = query(tree, node*2+1, mid+1, end, left, right);
    
    return leftChild + rightChild;
}

void update(vector<ll> &tree, int node, int start, int end, int index, ll diff) {
    // case 2: [start, end]가 [left, right]에 포함
    if (index < start || index > end) return;

    // case 1: [start, end] 앞 뒤에 [left, right]가 있는 경우
    tree[node] = tree[node] + diff;

    // 리프 노드가 아닌 경우에는 자식 노드도 변경해주어야 함
    if (start != end) {
        int mid = (start + end) / 2;
        update(tree, node*2, start, mid, index, diff);
        update(tree, node*2+1, mid+1, end, index, diff);
    }
}

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

    int n, m, k; // 수의 개수, update 횟수, query 횟수
    cin >> n >> m >> k;

    vector<ll> arr = vector<ll>(n);
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
    }
    vector<ll> tree = vector<ll>(4*n);
    init(arr, tree, 1, 0, n-1);

    // a = 1: b번째 수를 c로 바꿔라
    // a = 2: b번째 수부터 c번째 수까지의 합을 구해라
    int a, b, c;
    for (int i = 0; i < m+k; ++i) {
        cin >> a >> b >> c;
        // 배열의 index는 0부터, b는 1부터 시작하므로 b를 1 빼준다
        b--;

        if (a == 1) update(tree, 1, 0, n-1, b, c - arr[b]);
        else cout << query(tree, 1, 0, n-1, b, c) << "\n";
    }

    return 0;
}

 

3. 구현 아이디어 2 - 맞았습니다!!

하지만 위의 코드는 2%쯤부터 틀렸습니다가 나왔다 :(

질문 게시판을 통해 내 코드에서 잘못된 점 2가지를 찾을 수 있었다.

 

이 2가지 실수를 수정해주어 문제를 최종적으로 맞출 수 있었다.

1. 원소를 변경할 때 arr에도 반영해줄 것
2. a=1일 때 c는 index가 아닌 입력값을 의미하므로 int가 아닌 long long으로 변경해줄 것
(코드에서 변경된 부분은 변경 포인트 1, 2라고 주석으로 표기해두었다.)

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

    int n, m, k; // 수의 개수, update 횟수, query 횟수
    cin >> n >> m >> k;

    vector<ll> arr = vector<ll>(n);
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
    }
    vector<ll> tree = vector<ll>(4*n);
    init(arr, tree, 1, 0, n-1);

    // a = 1: b번째 수를 c로 바꿔라
    // a = 2: b번째 수부터 c번째 수까지의 합을 구해라
    int a, b;
    
    // 변경 포인트 1) 입력은 최대 2^63까지이므로 long long으로 변경
    ll c;
    for (int i = 0; i < m+k; ++i) {
        cin >> a >> b >> c;
        b--;

        if (a == 1) {
            update(tree, 1, 0, n-1, b, c - arr[b]);
            
            // 변경 포인트 2) arr에도 원소가 변경된 것을 반영
            arr[b] = c;
        }
        else cout << query(tree, 1, 0, n-1, b, c-1) << "\n";
    }

    return 0;
}